feat(cwb): save merge metadata
This commit is contained in:
@@ -241,6 +241,12 @@ alignment, weighting, or merge. Genuine factor ranks remain variable during
|
||||
CWB. The first maximum-rank input supplies the output component space, so the
|
||||
output rank remains the largest genuine input rank. Merged outputs never
|
||||
contain or invent alpha tensors.
|
||||
|
||||
Every CWB merge node writes a `cwb.merge` safetensors metadata record. It
|
||||
contains selected source model names, fully resolved CWB settings, and merge
|
||||
options. Presets are stored as effective settings rather than by preset name.
|
||||
Non-empty exclude/discard filters and their glob mode are included. CWB reports
|
||||
and diagnostics are not saved in metadata.
|
||||
Similarity alignment operates on paired latent-rank components. A rows and B
|
||||
columns share one mapping; fixed input/output feature axes are never reordered.
|
||||
For a rank-R reference and rank-r source it evaluates an R-by-r similarity
|
||||
|
||||
@@ -2,10 +2,11 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from typing import Iterable
|
||||
|
||||
import comfy.utils
|
||||
@@ -796,6 +797,51 @@ def build_custom_cwb_settings(**values) -> CWBSettings:
|
||||
return settings
|
||||
|
||||
|
||||
def _cwb_merge_metadata(
|
||||
source_models: list[str],
|
||||
params: dict,
|
||||
settings: CWBSettings,
|
||||
*,
|
||||
operation: str,
|
||||
) -> str:
|
||||
metadata = {
|
||||
"schema_version": 1,
|
||||
"operation": operation,
|
||||
"source_models": source_models,
|
||||
"effective_cwb_settings": asdict(settings),
|
||||
"merge_options": {},
|
||||
}
|
||||
for key in (
|
||||
"mismatch_mode",
|
||||
"save_dtype",
|
||||
"process_device",
|
||||
"lazy_load",
|
||||
"force_clear_cache",
|
||||
"override_dtype",
|
||||
"include_1d_diffs",
|
||||
"counterfactual_weight_sweep",
|
||||
"embedding_coalesce",
|
||||
"target_vector_count",
|
||||
"coalesce_similarity_threshold",
|
||||
"coalesce_position_window",
|
||||
"vision_boundary_embeddings",
|
||||
"legacy_boundary_search",
|
||||
"boundary_reference_embedding",
|
||||
"boundary_similarity_threshold",
|
||||
):
|
||||
if key in params:
|
||||
metadata["merge_options"][key] = params[key]
|
||||
exclude_patterns = params["exclude_patterns"].split()
|
||||
discard_patterns = params["discard_patterns"].split()
|
||||
if exclude_patterns or discard_patterns:
|
||||
metadata["filters"] = {
|
||||
"exclude_patterns": exclude_patterns,
|
||||
"discard_patterns": discard_patterns,
|
||||
"glob_patterns": params["glob_patterns"],
|
||||
}
|
||||
return json.dumps(metadata, sort_keys=True, separators=(",", ":"))
|
||||
|
||||
|
||||
def _common_prefix_length(tensors: list[torch.Tensor]) -> int:
|
||||
if not tensors:
|
||||
return 0
|
||||
@@ -1622,7 +1668,7 @@ class ConsensusMergerLogic:
|
||||
)
|
||||
if lora_mode:
|
||||
result = cls._merge_loras(
|
||||
handlers, low_bit_sets, model_type, params, settings
|
||||
handlers, low_bit_sets, model_type, params, settings, model_names
|
||||
)
|
||||
else:
|
||||
result = cls._merge_generic(
|
||||
@@ -1633,6 +1679,7 @@ class ConsensusMergerLogic:
|
||||
settings,
|
||||
embedding_union=embedding_union,
|
||||
boundary_handler=boundary_handler,
|
||||
source_models=model_names,
|
||||
)
|
||||
if params.get("_return_cwb_report"):
|
||||
report = diagnostics.render(
|
||||
@@ -1672,6 +1719,7 @@ class ConsensusMergerLogic:
|
||||
*,
|
||||
embedding_union,
|
||||
boundary_handler=None,
|
||||
source_models,
|
||||
):
|
||||
primary = handlers[0]
|
||||
keys = set()
|
||||
@@ -1699,7 +1747,14 @@ class ConsensusMergerLogic:
|
||||
}
|
||||
work_units.append((key, entries))
|
||||
|
||||
with atomic_uel_writer(output_path, primary.metadata()) as writer, torch.no_grad(), closing(
|
||||
output_metadata = (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",
|
||||
)
|
||||
with atomic_uel_writer(output_path, output_metadata) as writer, torch.no_grad(), closing(
|
||||
stream_work_units(
|
||||
handler_map, work_units,
|
||||
pin_memory=str(params["process_device"]).startswith("cuda"),
|
||||
@@ -1917,7 +1972,9 @@ class ConsensusMergerLogic:
|
||||
return cls._output_name(model_type, output_path)
|
||||
|
||||
@classmethod
|
||||
def _merge_loras(cls, handlers, low_bit_sets, model_type, params, settings):
|
||||
def _merge_loras(
|
||||
cls, handlers, low_bit_sets, model_type, params, settings, source_loras
|
||||
):
|
||||
parsed = [parse_lora_layers(handler.keys()) for handler in handlers]
|
||||
logical_maps = []
|
||||
logical_cores = []
|
||||
@@ -2077,6 +2134,9 @@ class ConsensusMergerLogic:
|
||||
work_units.append((("passthrough", key), entries))
|
||||
|
||||
output_metadata = (handlers[0].metadata() or {}).copy()
|
||||
output_metadata["cwb.merge"] = _cwb_merge_metadata(
|
||||
source_loras, params, settings, operation="cwb_lora_merge"
|
||||
)
|
||||
output_metadata["alpha_normalized"] = "true"
|
||||
output_metadata["alpha_normalization"] = (
|
||||
"lora_up := lora_up * (alpha / rank); alpha tensors removed"
|
||||
|
||||
Reference in New Issue
Block a user