From 8286c2d5d99c3d102944796e11638790621e5b67 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 21 Oct 2026 08:34:24 +0200 Subject: [PATCH] Add standalone CWBModelGraft node --- nodes/consensus_merger.py | 743 +++++++++++++++++++++++++++++++++ tests/test_consensus_merger.py | 133 +++++- 2 files changed, 874 insertions(+), 2 deletions(-) diff --git a/nodes/consensus_merger.py b/nodes/consensus_merger.py index c3a2921..f340f56 100644 --- a/nodes/consensus_merger.py +++ b/nodes/consensus_merger.py @@ -5,6 +5,7 @@ from __future__ import annotations import json import logging import os +import re from contextlib import closing from dataclasses import asdict, dataclass, field from typing import Any, Iterable @@ -3353,6 +3354,747 @@ class CWBEmbeddingStackMerger(CWBEmbeddingMultiMerger): return io.NodeOutput(filename, documentation, report) +_MAIN_BLOCK_INDEX_RE = re.compile( + r"(?:^|\.)(?:layers|blocks|double_blocks|single_blocks|transformer_blocks)\.(\d+)\." +) +_ATTN_KEY_TOKENS = ( + ".qkv.", + ".qkv_proj.", + ".proj.q.", + ".proj.k.", + ".proj.v.", + ".proj.o.", + ".attn.", + ".self_attn.", + ".q_norm.", + ".k_norm.", + ".to_q.", + ".to_k.", + ".to_v.", + ".to_out.", +) +_FFN_KEY_TOKENS = ( + ".ffn_up_gate.", + ".ffn_up.", + ".ffn_gate.", + ".ffn_down.", + ".mlp.", + ".ff.", + ".fc1.", + ".fc2.", + ".w1.", + ".w2.", + ".w3.", +) + + +def _extract_main_block_index(key: str) -> int | None: + if "text_embedder." in key or "token_refiner." in key: + return None + match = _MAIN_BLOCK_INDEX_RE.search(key) + return int(match.group(1)) if match else None + + +def _classify_graft_operator(key: str) -> str | None: + padded = f".{key}." + if any(token in padded for token in _ATTN_KEY_TOKENS): + return "attention" + if any(token in padded for token in _FFN_KEY_TOKENS): + return "ffn" + return None + + +def _block_rule_allows( + block_idx: int | None, + rule: str, + protect_early_blocks: int, + max_block_index: int, +) -> bool: + if rule == "none": + return False + if block_idx is None: + return False + if block_idx < protect_early_blocks: + return False + if rule == "all_after_protected": + return True + if rule == "interleaved_even": + return (block_idx % 2) == 0 + if rule == "interleaved_odd": + return (block_idx % 2) == 1 + if rule == "back_half_only": + midpoint = (max_block_index + 1) // 2 + return block_idx >= max(protect_early_blocks, midpoint) + return False + + +def _parse_custom_key_map(raw_map: str) -> dict[str, str]: + """Parse lines of 'donor_key -> base_key' or 'donor_key = base_key' into {base_key: donor_key}.""" + mapping: dict[str, str] = {} + if not raw_map or not raw_map.strip(): + return mapping + for line_num, raw_line in enumerate(raw_map.splitlines(), start=1): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if "->" in line: + donor_part, base_part = line.split("->", 1) + elif "=" in line: + donor_part, base_part = line.split("=", 1) + else: + raise ValueError( + f"Custom key map line {line_num}: expected 'donor_key -> base_key'." + ) + donor_key = donor_part.strip() + base_key = base_part.strip() + if not donor_key or not base_key: + raise ValueError( + f"Custom key map line {line_num}: both donor_key and base_key must be non-empty." + ) + mapping[base_key] = donor_key + return mapping + + +def _align_donor_trailing_shape( + base_tensor: torch.Tensor, donor_tensor: torch.Tensor +) -> tuple[torch.Tensor, bool]: + """Align donor trailing dimensions to base_tensor so CWB similarity alignment can match rows.""" + if donor_tensor.ndim != base_tensor.ndim: + if donor_tensor.numel() == base_tensor.numel(): + return donor_tensor.reshape(base_tensor.shape), True + flat_donor = donor_tensor.reshape(1, 1, -1).to(dtype=torch.float32) + resized = F.interpolate( + flat_donor, size=base_tensor.numel(), mode="linear", align_corners=False + ) + return resized.reshape(base_tensor.shape).to(dtype=donor_tensor.dtype), True + + if donor_tensor.shape == base_tensor.shape: + return donor_tensor, False + + if base_tensor.ndim <= 1: + return donor_tensor, True + + if donor_tensor.shape[1:] == base_tensor.shape[1:]: + return donor_tensor, True + + target_trailing = int(torch.tensor(base_tensor.shape[1:]).prod().item()) + donor_2d = donor_tensor.reshape(donor_tensor.shape[0], 1, -1).to(dtype=torch.float32) + resized_2d = F.interpolate( + donor_2d, size=target_trailing, mode="linear", align_corners=False + ) + reshaped = resized_2d.reshape(donor_tensor.shape[0], *base_tensor.shape[1:]).to( + dtype=donor_tensor.dtype + ) + return reshaped, True + + +def _graft_pair_to_cpu( + base_tensor: torch.Tensor, + donor_tensor: torch.Tensor, + settings: CWBSettings, + device: str, + target_dtype: torch.dtype, + *, + operation_label: str, + diagnostics: CWBDiagnostics, +) -> torch.Tensor: + aligned_donor, shape_changed = _align_donor_trailing_shape(base_tensor, donor_tensor) + merged = _merge_tensors_to_cpu( + [base_tensor, aligned_donor], + settings, + device, + torch.float32, + reference_index=0, + allow_similarity_alignment=True, + operation_label=operation_label, + diagnostics=diagnostics, + ) + if shape_changed or settings.rescale_norm: + base_norm = torch.linalg.vector_norm(base_tensor.to(dtype=torch.float32)) + merged_norm = torch.linalg.vector_norm(merged) + if merged_norm.item() > 1e-8 and base_norm.item() > 0.0: + merged = merged * (base_norm / merged_norm) + return merged.to(dtype=target_dtype).cpu().contiguous() + + +def _graft_with_subpart( + key: str, + op_type: str, + subpart: str, + base_tensor: torch.Tensor, + donor_tensor: torch.Tensor, + settings: CWBSettings, + device: str, + target_dtype: torch.dtype, + diagnostics: CWBDiagnostics, +) -> torch.Tensor | None: + """Return grafted tensor or None if the subpart rule keeps the base tensor untouched.""" + padded = f".{key}." + if op_type == "attention": + is_qkv = any(tok in padded for tok in (".qkv.", ".qkv_proj.")) + is_q_or_k = any( + tok in padded for tok in (".proj.q.", ".proj.k.", ".to_q.", ".to_k.", ".q_norm.", ".k_norm.") + ) + is_v = any(tok in padded for tok in (".proj.v.", ".to_v.")) + is_out = any(tok in padded for tok in (".proj.o.", ".out_proj.", ".to_out.")) + has_known_attn_part = is_qkv or is_q_or_k or is_v or is_out + + if has_known_attn_part and subpart == "output_only": + if not is_out: + return None + elif has_known_attn_part and subpart == "value_and_output": + if is_q_or_k: + return None + if is_qkv and base_tensor.ndim >= 1 and base_tensor.shape[0] % 3 == 0: + aligned_donor, _ = _align_donor_trailing_shape(base_tensor, donor_tensor) + if aligned_donor.shape[0] % 3 == 0: + q_a, k_a, v_a = base_tensor.chunk(3, dim=0) + _, _, v_b = aligned_donor.chunk(3, dim=0) + v_merged = _graft_pair_to_cpu( + v_a, + v_b, + settings, + device, + target_dtype, + operation_label=f"{key}[v]", + diagnostics=diagnostics, + ) + return torch.cat( + [ + q_a.to(dtype=target_dtype).cpu(), + k_a.to(dtype=target_dtype).cpu(), + v_merged, + ], + dim=0, + ).contiguous() + elif not (is_v or is_out): + return None + + elif op_type == "ffn": + is_up_gate = ".ffn_up_gate." in padded + is_up = any(tok in padded for tok in (".ffn_up.", ".fc1.", ".w1.")) + is_gate = any(tok in padded for tok in (".ffn_gate.", ".w3.")) + is_down = any(tok in padded for tok in (".ffn_down.", ".fc2.", ".w2.", ".ff.net.2.")) + has_known_ffn_part = is_up_gate or is_up or is_gate or is_down + + if has_known_ffn_part and subpart == "down_only": + if not is_down: + return None + elif has_known_ffn_part and subpart == "up_only": + if is_down or is_gate: + return None + if is_up_gate and base_tensor.ndim >= 1 and base_tensor.shape[0] % 2 == 0: + aligned_donor, _ = _align_donor_trailing_shape(base_tensor, donor_tensor) + if aligned_donor.shape[0] % 2 == 0: + up_a, gate_a = base_tensor.chunk(2, dim=0) + up_b, _ = aligned_donor.chunk(2, dim=0) + up_merged = _graft_pair_to_cpu( + up_a, + up_b, + settings, + device, + target_dtype, + operation_label=f"{key}[up]", + diagnostics=diagnostics, + ) + return torch.cat( + [ + up_merged, + gate_a.to(dtype=target_dtype).cpu(), + ], + dim=0, + ).contiguous() + elif not is_up: + return None + + return _graft_pair_to_cpu( + base_tensor, + donor_tensor, + settings, + device, + target_dtype, + operation_label=key, + diagnostics=diagnostics, + ) + + +def _resolve_donor_qkv_plan( + base_key: str, + donor_keys_set: set[str], + custom_map: dict[str, str], +) -> tuple[str, list[str], int | None] | None: + """Resolve donor key(s) for base_key. + + Returns: + (mode, donor_logical_keys, slice_index) + mode is 'direct', 'fuse_qkv' (cat q,k,v), or 'split_qkv' (chunk 3 and pick slice_index). + """ + if base_key in custom_map: + mapped = custom_map[base_key] + if mapped in donor_keys_set: + return ("direct", [mapped], None) + return None + if base_key in donor_keys_set: + return ("direct", [base_key], None) + + if base_key.endswith(".proj.qkv.weight"): + prefix = base_key[: -len("qkv.weight")] + parts = [f"{prefix}{p}.weight" for p in ("q", "k", "v")] + if all(k in donor_keys_set for k in parts): + return ("fuse_qkv", parts, None) + + for idx, part in enumerate(("q", "k", "v")): + suffix = f".proj.{part}.weight" + if base_key.endswith(suffix): + fused = f"{base_key[:-len(suffix)]}.proj.qkv.weight" + if fused in donor_keys_set: + return ("split_qkv", [fused], idx) + + return None + + +class CWBModelGraft(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + models = folder_paths.get_filename_list("diffusion_models") + presets = list(DENSE_CWB_PRESETS) + return io.Schema( + node_id="CWBModelGraft", + display_name="CWB Model Graft", + category="ModelUtils/Merging", + description="Graft selected blocks, sub-projections, or mapped layers from a donor diffusion model into a base diffusion model using Consensus-Weighted Blending.", + inputs=[ + io.Combo.Input( + "execution_mode", + options=["MERGE", "DOCUMENTATION ONLY"], + default="MERGE", + tooltip="MERGE writes a grafted model; DOCUMENTATION ONLY opens no model files.", + ), + io.Combo.Input( + "model_a", + options=models, + tooltip="Base diffusion model that anchors structure, metadata, and protected blocks.", + ), + io.Combo.Input( + "model_b", + options=models, + tooltip="Donor diffusion model supplying parts to graft into Model A.", + ), + io.Combo.Input( + "attention_preset", + options=presets, + default="robust_medn", + tooltip="CWB preset applied to grafted attention layers.", + ), + CWB_CONFIG.Input( + "attention_config", + optional=True, + tooltip="Optional custom CWB configuration override for attention layers.", + ), + io.Combo.Input( + "ffn_preset", + options=presets, + default="balanced_mean", + tooltip="CWB preset applied to grafted feed-forward (MLP) layers.", + ), + CWB_CONFIG.Input( + "ffn_config", + optional=True, + tooltip="Optional custom CWB configuration override for feed-forward (MLP) layers.", + ), + io.Int.Input( + "protect_early_blocks", + default=2, + min=0, + max=256, + step=1, + tooltip="Number of initial transformer blocks (starting from 0) kept 100% from Model A to preserve global layout.", + ), + io.Combo.Input( + "attention_block_rule", + options=[ + "interleaved_even", + "interleaved_odd", + "back_half_only", + "all_after_protected", + "none", + ], + default="interleaved_even", + tooltip="Which transformer blocks after the protected start have their attention layers grafted.", + ), + io.Combo.Input( + "ffn_block_rule", + options=[ + "all_after_protected", + "interleaved_even", + "interleaved_odd", + "back_half_only", + "none", + ], + default="all_after_protected", + tooltip="Which transformer blocks after the protected start have their feed-forward layers grafted.", + ), + io.Combo.Input( + "attention_subpart", + options=["all", "value_and_output", "output_only"], + default="all", + tooltip="Select whether to graft all attention weights, only value and output projections, or only the output projection.", + ), + io.Combo.Input( + "ffn_subpart", + options=["all", "up_only", "down_only"], + default="all", + tooltip="Select whether to graft all feed-forward weights, only the up projection, or only the down projection.", + ), + io.String.Input( + "custom_key_map", + default="", + multiline=True, + tooltip="Optional explicit key mappings (one per line as 'donor_key -> base_key') to graft across differently named layers.", + ), + io.String.Input( + "output_filename", + default="cwb_grafted_model", + tooltip="Filename without extension under ComfyUI's diffusion_models directory.", + ), + io.Combo.Input( + "save_dtype", + options=["fp32", "fp16", "bf16"], + default="bf16", + tooltip="Requested dtype for grafted floating tensors.", + ), + io.Combo.Input( + "process_device", + options=["cuda", "cpu"], + default="cuda", + tooltip="Per-tensor FP32 CWB computation device; CUDA OOM retries on CPU.", + ), + io.String.Input( + "exclude_patterns", + default="", + multiline=True, + tooltip="Newline or whitespace separated patterns to preserve from Model A (or whitelist when include_mode is enabled).", + ), + io.String.Input( + "discard_patterns", + default="", + multiline=True, + tooltip="Newline or whitespace separated patterns to omit from the output file.", + ), + io.Boolean.Input( + "glob_patterns", + default=False, + tooltip="Interpret exclude and discard entries as shell-style globs instead of regular expressions.", + ), + io.Boolean.Input( + "lazy_load", + default=True, + tooltip="Stream only one logical work unit at a time through UEL and release completed tensors immediately.", + ), + io.Boolean.Input( + "force_clear_cache", + default=False, + tooltip="Run Python garbage collection and empty CUDA cache before each tensor.", + ), + io.Boolean.Input( + "override_dtype", + default=False, + tooltip="Force grafted floating tensors to Save Dtype instead of preserving Model A's float precision.", + ), + io.Boolean.Input( + "include_mode", + default=False, + tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are eligible for grafting; nonmatching layers are preserved from Model A.", + ), + ], + outputs=[ + io.AnyType.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + io.String.Output(display_name="cwb_report"), + ], + is_experimental=True, + ) + + @classmethod + def execute(cls, **kwargs) -> io.NodeOutput: + documentation = load_documentation_from_file("consensus_mergers.md") + if kwargs["execution_mode"] == "DOCUMENTATION ONLY": + return io.NodeOutput( + "Documentation mode active. No merge performed.", + documentation, + "CWB report unavailable because no merge was performed.", + ) + filename, report = cls._run_graft(kwargs) + return io.NodeOutput(filename, documentation, report) + + @classmethod + def _run_graft(cls, params: dict) -> tuple[str, str]: + model_names = [params["model_a"], params["model_b"]] + paths = [] + for name in model_names: + path = folder_paths.get_full_path("diffusion_models", name) + if not path: + raise FileNotFoundError(f"diffusion_models 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)) + + custom_map = _parse_custom_key_map(params.get("custom_key_map", "")) + attn_settings = resolve_cwb_settings( + params["attention_preset"], + DENSE_CWB_PRESETS, + params.get("attention_config"), + ) + ffn_settings = resolve_cwb_settings( + params["ffn_preset"], + DENSE_CWB_PRESETS, + params.get("ffn_config"), + ) + diagnostics = CWBDiagnostics() + protect_early = int(params.get("protect_early_blocks", 2)) + attn_rule = params.get("attention_block_rule", "interleaved_even") + ffn_rule = params.get("ffn_block_rule", "all_after_protected") + attn_subpart = params.get("attention_subpart", "all") + ffn_subpart = params.get("ffn_subpart", "all") + glob_mode = bool(params.get("glob_patterns", False)) + include_mode = bool(params.get("include_mode", False)) + exclude = _compile_patterns(params.get("exclude_patterns", ""), glob_mode=glob_mode) + discard = _compile_patterns(params.get("discard_patterns", ""), glob_mode=glob_mode) + requested_dtype = _requested_dtype(params["save_dtype"]) + + handlers = [ + MemoryEfficientSafeOpen(path, low_memory=params["lazy_load"]) + for path in paths + ] + try: + quantizers = [ + DiffusionQuantization(handler, path, f"CWB Graft input {idx + 1}") + for idx, (handler, path) in enumerate(zip(handlers, paths)) + ] + key_maps = [ + diffusion_key_map(q.data_keys, f"CWB Graft input {idx + 1}") + for idx, q in enumerate(quantizers) + ] + low_bit_sets = [ + { + logical + for logical, physical in mapping.items() + if physical in q.isolated_low_bit_keys + } + for mapping, q in zip(key_maps, quantizers) + ] + + base_keys = sorted(key_maps[0]) + donor_keys_set = set(key_maps[1]) + block_indices = [ + idx + for k in base_keys + if (idx := _extract_main_block_index(k)) is not None + ] + max_block_index = max(block_indices) if block_indices else 0 + + def matches_filter(key: str, patterns: list) -> bool: + aliases = [key, key_maps[0][key]] + return any( + _matches_any_pattern(alias, patterns, glob_mode=glob_mode) + for alias in aliases + ) + + def should_graft_key(key: str) -> tuple[bool, str | None, tuple[str, list[str], int | None] | None]: + if matches_filter(key, discard): + return (False, None, None) + has_filter_rules = bool(exclude) + matched_filter = matches_filter(key, exclude) + if include_mode and not matched_filter: + return (False, None, None) + if not include_mode and matched_filter: + return (False, None, None) + + donor_plan = _resolve_donor_qkv_plan(key, donor_keys_set, custom_map) + if donor_plan is None: + return (False, None, None) + _, donor_logical_keys, _ = donor_plan + if key in low_bit_sets[0] or any(dk in low_bit_sets[1] for dk in donor_logical_keys): + return (False, None, None) + + op_type = _classify_graft_operator(key) + if key in custom_map or (include_mode and has_filter_rules and matched_filter): + return (True, op_type or "ffn", donor_plan) + + block_idx = _extract_main_block_index(key) + if op_type == "attention": + if _block_rule_allows(block_idx, attn_rule, protect_early, max_block_index): + return (True, "attention", donor_plan) + elif op_type == "ffn": + if _block_rule_allows(block_idx, ffn_rule, protect_early, max_block_index): + return (True, "ffn", donor_plan) + return (False, None, None) + + work_units = [] + graft_plans: dict[str, tuple[str, tuple[str, list[str], int | None]]] = {} + for key in base_keys: + if matches_filter(key, discard): + work_units.append((key, {})) + continue + graft_ok, op_type, donor_plan = should_graft_key(key) + base_phys = key_maps[0][key] + if not graft_ok: + entries = ( + {0: quantizers[0].required_keys(base_phys)} + if base_phys in quantizers[0].quantized_keys + else {} + ) + work_units.append((key, entries)) + continue + + assert op_type is not None and donor_plan is not None + graft_plans[key] = (op_type, donor_plan) + donor_req: list[str] = [] + for dk in donor_plan[1]: + d_phys = key_maps[1][dk] + for rk in quantizers[1].required_keys(d_phys): + if rk not in donor_req: + donor_req.append(rk) + work_units.append( + ( + key, + { + 0: quantizers[0].required_keys(base_phys), + 1: donor_req, + }, + ) + ) + + output_path, output_rel = canonical_model_artifact_path( + "diffusion_models", params["output_filename"] + ) + output_metadata = quantizers[0].output_metadata() + output_metadata["cwb.graft"] = json.dumps( + { + "base_model": params["model_a"], + "donor_model": params["model_b"], + "attention_preset": params["attention_preset"], + "ffn_preset": params["ffn_preset"], + "protect_early_blocks": protect_early, + "attention_block_rule": attn_rule, + "ffn_block_rule": ffn_rule, + "attention_subpart": attn_subpart, + "ffn_subpart": ffn_subpart, + }, + sort_keys=True, + ) + + def decode_logical(index: int, logical_key: str, loaded: dict) -> torch.Tensor: + phys = key_maps[index][logical_key] + if phys in quantizers[index].quantized_keys: + parts = { + part: loaded[(index, part)] + for part in quantizers[index].required_keys(phys) + } + return quantizers[index].decode(phys, parts, torch.float32, device="cpu") + return loaded[(index, phys)] + + def preserve_base(writer, key: str, loaded: dict, *, force_raw: bool = False): + phys = key_maps[0][key] + if phys in quantizers[0].quantized_keys and (0, phys) in loaded: + decoded = decode_logical(0, key, loaded) + writer.write(key, decoded.to(requested_dtype).cpu().contiguous()) + else: + write_preserved_tensor( + writer, + phys, + handlers[0], + output_key=key, + force_raw=force_raw, + tensor=loaded.get((0, phys)), + ) + + handler_map = {0: handlers[0], 1: handlers[1]} + grafted_count = 0 + preserved_count = 0 + pbar = comfy.utils.ProgressBar(len(work_units)) + with atomic_uel_writer(output_path, output_metadata) as writer, torch.no_grad(), closing( + stream_work_units( + handler_map, + work_units, + pin_memory=str(device).startswith("cuda"), + ) + ) as streamed: + for key, loaded in tqdm(streamed, total=len(work_units), desc="CWB grafting tensors", unit="tensors"): + _clear_previous_layer(params) + if matches_filter(key, discard): + pbar.update(1) + continue + if key not in graft_plans: + preserved_count += 1 + preserve_base(writer, key, loaded, force_raw=key in low_bit_sets[0]) + pbar.update(1) + continue + + op_type, (mode, donor_logical_keys, slice_idx) = graft_plans[key] + base_tensor = decode_logical(0, key, loaded) + if mode == "direct": + donor_tensor = decode_logical(1, donor_logical_keys[0], loaded) + elif mode == "fuse_qkv": + donor_tensor = torch.cat( + [decode_logical(1, dk, loaded) for dk in donor_logical_keys], + dim=0, + ) + else: + fused_donor = decode_logical(1, donor_logical_keys[0], loaded) + donor_tensor = fused_donor.chunk(3, dim=0)[slice_idx or 0] + + base_phys = key_maps[0][key] + base_dtype = quantizers[0].logical_dtype(base_phys, requested_dtype) + if not _is_float_dtype(base_dtype) or not _is_float_dtype(donor_tensor.dtype): + preserved_count += 1 + preserve_base(writer, key, loaded) + pbar.update(1) + continue + + target_dtype = ( + requested_dtype + if params.get("override_dtype", False) + else base_dtype + ) + active_settings = attn_settings if op_type == "attention" else ffn_settings + active_subpart = attn_subpart if op_type == "attention" else ffn_subpart + grafted = _graft_with_subpart( + key, + op_type, + active_subpart, + base_tensor, + donor_tensor, + active_settings, + device, + target_dtype, + diagnostics, + ) + if grafted is None: + preserved_count += 1 + preserve_base(writer, key, loaded) + else: + grafted_count += 1 + writer.write(key, grafted) + del base_tensor, donor_tensor, grafted + pbar.update(1) + + report = diagnostics.render( + f"attn={params['attention_preset']}, ffn={params['ffn_preset']}", + 2, + params.get("attention_config") is not None or params.get("ffn_config") is not None, + ) + report += f"\nGrafted tensors: {grafted_count}\nPreserved base tensors: {preserved_count}\n" + return output_rel, report + finally: + for handler in handlers: + handler.__exit__(None, None, None) + cleanup_after_operation() + + CWB_MERGER_NODES = [ CWBCustomConfiguration, CWBEmbeddingStack, @@ -3360,6 +4102,7 @@ CWB_MERGER_NODES = [ CWBCheckpointThreeMerger, CWBModelTwoMerger, CWBModelThreeMerger, + CWBModelGraft, CWBTextEncoderTwoMerger, CWBTextEncoderThreeMerger, CWBLoRATwoMerger, diff --git a/tests/test_consensus_merger.py b/tests/test_consensus_merger.py index 836752d..ce34a17 100644 --- a/tests/test_consensus_merger.py +++ b/tests/test_consensus_merger.py @@ -622,9 +622,9 @@ def test_norm_rescale_and_dsc_bandpass_are_finite(cwb): def test_all_cwb_schemas_and_compact_control_contract(cwb): - assert len(cwb.CWB_MERGER_NODES) == 16 + assert len(cwb.CWB_MERGER_NODES) == 17 schemas = [node.define_schema() for node in cwb.CWB_MERGER_NODES] - assert len({schema.node_id for schema in schemas}) == 16 + assert len({schema.node_id for schema in schemas}) == 17 advanced = { "consensus_type", "alignment_method", "alignment_threshold", "similarity_threshold", "power_alpha", "diversity_beta", @@ -648,6 +648,12 @@ def test_all_cwb_schemas_and_compact_control_contract(cwb): assert schema.outputs[1].get_io_type() == "STRING" assert schema.outputs[2].get_io_type() == "STRING" assert advanced.isdisjoint(ids) + if node is cwb.CWBModelGraft: + assert "attention_preset" in ids + assert "ffn_preset" in ids + assert "protect_early_blocks" in ids + assert ids[-1] == "include_mode" + continue assert "cwb_config" in ids if node is cwb.CWBLoRAMultiMerger: assert ids[:4] == ["execution_mode", "lora_count", "lora_1", "lora_2"] @@ -1159,6 +1165,129 @@ def test_embedding_stack_and_stack_merger(monkeypatch, cwb): ) +def test_cwb_model_graft_protected_blocks_subparts_and_mismatched_shapes( + monkeypatch, tmp_path, cwb +): + base_path = tmp_path / "base.safetensors" + donor_path = tmp_path / "donor.safetensors" + + base_qkv_0 = torch.ones((6, 4), dtype=torch.float32) * 2.0 + base_qkv_2 = torch.cat( + [ + torch.ones((2, 4), dtype=torch.float32) * 1.0, + torch.ones((2, 4), dtype=torch.float32) * 2.0, + torch.ones((2, 4), dtype=torch.float32) * 3.0, + ], + dim=0, + ) + base_ffn_2 = torch.cat( + [ + torch.ones((2, 4), dtype=torch.float32) * 4.0, + torch.ones((2, 4), dtype=torch.float32) * 5.0, + ], + dim=0, + ) + base_custom = torch.ones((2, 4), dtype=torch.float32) * 2.0 + + save_file( + { + "layers.0.proj.qkv.weight": base_qkv_0, + "layers.2.proj.qkv.weight": base_qkv_2, + "layers.2.proj.ffn_up_gate.weight": base_ffn_2, + "layers.3.proj.qkv.weight": base_qkv_0.clone(), + "custom_target.weight": base_custom, + }, + str(base_path), + ) + + # Donor has split q/k/v for layers.2 and a differently-shaped custom_donor.weight + save_file( + { + "layers.0.proj.qkv.weight": torch.ones((6, 4), dtype=torch.float32) * 99.0, + "layers.2.proj.q.weight": torch.ones((2, 4), dtype=torch.float32) * 50.0, + "layers.2.proj.k.weight": torch.ones((2, 4), dtype=torch.float32) * 60.0, + "layers.2.proj.v.weight": torch.tensor( + [[9.0, 1.0, 1.0, 1.0], [1.0, 9.0, 1.0, 1.0]], dtype=torch.float32 + ), + "layers.2.proj.ffn_up_gate.weight": torch.cat( + [ + torch.tensor( + [[8.0, 1.0, 1.0, 1.0], [1.0, 8.0, 1.0, 1.0]], + dtype=torch.float32, + ), + torch.ones((2, 4), dtype=torch.float32) * 77.0, + ], + dim=0, + ), + "layers.3.proj.qkv.weight": torch.ones((6, 4), dtype=torch.float32) * 88.0, + "custom_donor.weight": torch.tensor( + [ + [10.0, 2.0, 1.0, 1.0, 1.0, 1.0], + [1.0, 10.0, 2.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 10.0, 2.0, 1.0, 1.0], + ], + dtype=torch.float32, + ), + }, + str(donor_path), + ) + + paths = {"base": str(base_path), "donor": str(donor_path)} + _patch_io(monkeypatch, cwb, tmp_path, paths) + monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS") + + out = cwb.CWBModelGraft.execute( + execution_mode="MERGE", + model_a="base", + model_b="donor", + attention_preset="balanced_mean", + attention_config=None, + ffn_preset="balanced_mean", + ffn_config=None, + protect_early_blocks=2, + attention_block_rule="interleaved_even", + ffn_block_rule="all_after_protected", + attention_subpart="value_and_output", + ffn_subpart="up_only", + custom_key_map="custom_donor.weight -> custom_target.weight", + output_filename="grafted", + save_dtype="fp32", + process_device="cpu", + exclude_patterns="", + discard_patterns="", + glob_patterns=False, + lazy_load=True, + force_clear_cache=False, + override_dtype=False, + include_mode=False, + ) + + rel_name, docs, report = out.result + assert docs == "DOCS" + assert "attn=balanced_mean, ffn=balanced_mean" in report + tensors = load_file(str(_result_path(tmp_path, "diffusion_models", rel_name))) + + # 1. Protected block 0 is byte-exact identical to Base + assert torch.equal(tensors["layers.0.proj.qkv.weight"], base_qkv_0) + # 2. Odd block 3 attention is untouched under interleaved_even + assert torch.equal(tensors["layers.3.proj.qkv.weight"], base_qkv_0) + # 3. Even block 2 qkv: Q and K slices stay Base, only V slice is grafted from split donor + q_out, k_out, v_out = tensors["layers.2.proj.qkv.weight"].chunk(3, dim=0) + assert torch.equal(q_out, base_qkv_2[:2]) + assert torch.equal(k_out, base_qkv_2[2:4]) + assert not torch.equal(v_out, base_qkv_2[4:6]) + # 4. Block 2 ffn_up_gate: up half is grafted, gate half stays Base + up_out, gate_out = tensors["layers.2.proj.ffn_up_gate.weight"].chunk(2, dim=0) + assert not torch.equal(up_out, base_ffn_2[:2]) + assert torch.equal(gate_out, base_ffn_2[2:]) + # 5. Mismatched-shape custom mapped tensor blends without changing shape or norm + assert tensors["custom_target.weight"].shape == base_custom.shape + assert not torch.equal(tensors["custom_target.weight"], base_custom) + assert torch.linalg.vector_norm(tensors["custom_target.weight"]).item() == pytest.approx( + torch.linalg.vector_norm(base_custom).item(), rel=1e-5 + ) + + def test_streaming_dtype_nonfloat_and_secondary_preservation( monkeypatch, tmp_path, cwb ):