nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(embeddings): add CWB coalescing

Browse Source
This commit is contained in:
silveroxides
2026-08-20 19:34:21 +02:00
parent d9467c4bf2
commit 066c7a34db
3 changed files with 509 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/consensus_mergers.md
+15
View File
@@ -17,6 +17,21 @@ Two-input and three-input variants are provided for checkpoints, standalone
diffusion models, text encoders, LoRAs, and embeddings. LoRA also provides a
Multi-Merge variant for 2 to 8 equal-prior inputs.
## Embedding coalescing
`CWB Embedding Self-Coalesce` reduces one embedding by replacing mutually
nearest body-vector pairs that meet both cosine and normalized-position limits.
`CWB Embedding Multi-Merge` accepts 2 to 8 embeddings, performs normal CWB
alignment first, then applies the same reduction. A target row count of zero
continues until no eligible pair remains.
Vision Boundary Embeddings preserves encoded vision-start and vision-end rows
unchanged and applies CWB only to their interior. Normal vision inputs validate
their outer rows before writing. Legacy Boundary Search trims accidental outer
template rows by matching a selected known-clean visual embedding; an absent,
ambiguous, or low-similarity boundary pair is an error rather than a guessed
crop. Disable Vision Boundary Embeddings for open textual-inversion tensors.
Model A supplies output metadata and anchors shared-layer shape, naming, and
mismatch behavior. The output key set is the union of all inputs. A valid
tensor found only in a secondary input is copied unchanged. When Model A lacks
nodes/consensus_merger.py
+361 -1
View File
@@ -255,6 +255,45 @@ class CWBSettings:
preserve_common_prefix: bool
@dataclass
class EmbeddingCoalesceReport:
"""Per-output accounting for dedicated embedding coalescing nodes."""
target_vector_count: int
similarity_threshold: float
position_window: float
vision_boundary_embeddings: bool
legacy_boundary_search: bool
boundary_similarity_threshold: float
boundary_reference: str | None
tensors: list[tuple[str, int, int, int, tuple[float, ...], tuple[tuple[int, int], ...]]] = field(default_factory=list)
def record(self, key, input_rows, output_rows, trimmed_rows, boundary_scores, pairs):
self.tensors.append((
key, input_rows, output_rows, trimmed_rows,
tuple(float(score) for score in boundary_scores), tuple(pairs),
))
def render(self) -> str:
lines = [
"", "CWB EMBEDDING COALESCING",
f"Target vector count: {self.target_vector_count} (0 means all eligible pairs)",
"Minimum retained body rows: 1",
f"Similarity threshold: {self.similarity_threshold:.4f}",
f"Position window: {self.position_window:.4f} normalized sequence distance",
f"Vision boundaries: {self.vision_boundary_embeddings}",
f"Legacy boundary search: {self.legacy_boundary_search}",
f"Boundary similarity threshold: {self.boundary_similarity_threshold:.4f}",
f"Boundary reference: {self.boundary_reference or 'none'}",
]
for key, before, after, trimmed, scores, pairs in self.tensors:
lines.append(
f"{key}: {before} -> {after} rows; trimmed={trimmed}; "
f"boundary_scores={[round(score, 6) for score in scores]}; pairs={list(pairs)}"
)
return "\n".join(lines)
def _optional_min(current: float | None, value: float | None) -> float | None:
if value is None:
return current
@@ -1183,6 +1222,136 @@ def _copy_to_target_shape(tensor: torch.Tensor, target: torch.Size) -> torch.Ten
return output
def _row_cosine(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
return F.cosine_similarity(left.reshape(1, -1), right.reshape(1, -1), dim=1)[0]
def _find_legacy_boundary_pair(
tensor: torch.Tensor,
start_reference: torch.Tensor,
end_reference: torch.Tensor,
threshold: float,
) -> tuple[int, int, tuple[float, float]]:
"""Locate one ordered encoded visual-boundary pair in a legacy tensor."""
if tensor.ndim != 2 or tensor.shape[0] < 2:
raise ValueError("Legacy vision boundary search requires a 2D tensor with at least two rows.")
rows = F.normalize(tensor.float(), p=2, dim=1, eps=1e-8)
start_scores = torch.mv(rows, F.normalize(start_reference.float(), p=2, dim=0, eps=1e-8))
end_scores = torch.mv(rows, F.normalize(end_reference.float(), p=2, dim=0, eps=1e-8))
pair_scores = start_scores[:, None] + end_scores[None, :]
ordered = torch.ones_like(pair_scores, dtype=torch.bool).triu_(diagonal=1)
pair_scores.masked_fill_(~ordered, -float("inf"))
flat_index = torch.argmax(pair_scores)
start = int((flat_index // tensor.shape[0]).item())
end = int((flat_index % tensor.shape[0]).item())
start_score = float(start_scores[start].item())
end_score = float(end_scores[end].item())
if start >= end or min(start_score, end_score) < threshold:
raise ValueError(
"Legacy vision boundary search found no ordered start/end pair meeting "
f"the boundary similarity threshold {threshold:.4f}."
)
return start, end, (start_score, end_score)
def _prepare_vision_embedding(
tensor: torch.Tensor,
*,
start_reference: torch.Tensor,
end_reference: torch.Tensor,
legacy_search: bool,
threshold: float,
) -> tuple[torch.Tensor, tuple[float, float], int]:
if tensor.ndim != 2 or tensor.shape[0] < 2:
raise ValueError("Vision boundary embeddings require a 2D tensor with start and end rows.")
if tensor.shape[1] != start_reference.numel() or tensor.shape[1] != end_reference.numel():
raise ValueError("Vision boundary embedding hidden sizes must match the selected boundary reference.")
if legacy_search:
start, end, scores = _find_legacy_boundary_pair(
tensor, start_reference, end_reference, threshold
)
else:
start, end = 0, tensor.shape[0] - 1
scores = (
float(_row_cosine(tensor[start], start_reference).item()),
float(_row_cosine(tensor[end], end_reference).item()),
)
if min(scores) < threshold:
raise ValueError(
"Vision boundary validation failed: encoded start/end rows do not meet "
f"the boundary similarity threshold {threshold:.4f}."
)
return tensor[start + 1:end], scores, start + (tensor.shape[0] - end - 1)
def coalesce_embedding_rows(
tensor: torch.Tensor,
settings: CWBSettings,
*,
target_vector_count: int,
similarity_threshold: float,
position_window: float,
) -> tuple[torch.Tensor, tuple[tuple[int, int], ...]]:
"""Repeatedly replace a mutual-nearest local pair with one CWB vector."""
if tensor.ndim != 2:
return tensor, ()
if target_vector_count < 0:
raise ValueError("Target vector count must be zero or positive.")
if not 0.0 <= similarity_threshold <= 1.0:
raise ValueError("Embedding coalescing similarity threshold must be between 0 and 1.")
if not 0.0 <= position_window <= 1.0:
raise ValueError("Embedding coalescing position window must be between 0 and 1.")
rows = tensor
merged_pairs = []
while rows.shape[0] > max(target_vector_count, 1):
count = rows.shape[0]
if count < 2:
break
normalized = F.normalize(rows.float(), p=2, dim=1, eps=1e-8)
scores = normalized @ normalized.t()
scores.fill_diagonal_(-float("inf"))
positions = torch.linspace(0.0, 1.0, count, device=rows.device)
scores[(positions[:, None] - positions[None, :]).abs() > position_window] = -float("inf")
nearest = scores.argmax(dim=1)
candidate_scores = scores[torch.arange(count, device=rows.device), nearest]
mutual = torch.arange(count, device=rows.device) == nearest[nearest]
candidate_scores[~mutual] = -float("inf")
candidate_scores[candidate_scores < similarity_threshold] = -float("inf")
first = int(torch.argmax(candidate_scores).item())
if not torch.isfinite(candidate_scores[first]):
break
second = int(nearest[first].item())
first, second = sorted((first, second))
merged = merge_consensus_group(torch.stack((rows[first], rows[second])), settings)
insert_at = (first + second) // 2
keep = torch.ones(count, dtype=torch.bool, device=rows.device)
keep[first] = False
keep[second] = False
remaining = rows[keep]
rows = torch.cat((remaining[:insert_at], merged.unsqueeze(0), remaining[insert_at:]), dim=0)
merged_pairs.append((first, second))
return rows, tuple(merged_pairs)
def _read_visual_boundary_reference(handler, key: str) -> tuple[torch.Tensor, torch.Tensor]:
"""Read and release only one reference tensor work unit through UEL."""
if key not in handler.keys():
raise ValueError(f"Boundary reference is missing embedding tensor '{key}'.")
stream = handler.async_stream([key], batch_size=1, prefetch_batches=1, pin_memory=False)
try:
batch = next(stream)
loaded_key, tensor = batch[0]
if loaded_key != key or tensor.ndim != 2 or tensor.shape[0] < 2:
raise ValueError("Boundary reference must contain a 2D embedding tensor with start and end rows.")
return tensor[0].detach().clone(), tensor[-1].detach().clone()
finally:
handler.mark_processed(key)
close = getattr(stream, "close", None)
if close is not None:
close()
def _to_compute(tensor: torch.Tensor, device: str) -> torch.Tensor:
if device == "cuda":
return transfer_to_gpu_pinned(tensor, device, torch.float32)
@@ -1425,6 +1594,15 @@ class ConsensusMergerLogic:
prepare_for_large_operation(total_size * 1.2, torch.device(device))
handlers = [MemoryEfficientSafeOpen(path, low_memory=params["lazy_load"]) for path in paths]
boundary_handler = None
boundary_name = params.get("boundary_reference_embedding")
if boundary_name and boundary_name != "None":
boundary_path = folder_paths.get_full_path("embeddings", boundary_name)
if not boundary_path:
raise FileNotFoundError(f"Boundary reference embedding '{boundary_name}' was not found.")
boundary_handler = MemoryEfficientSafeOpen(
boundary_path, low_memory=params["lazy_load"]
)
try:
low_bit_sets = [
inspect_low_bit_input(
@@ -1454,17 +1632,24 @@ class ConsensusMergerLogic:
params,
settings,
embedding_union=embedding_union,
boundary_handler=boundary_handler,
)
if params.get("_return_cwb_report"):
return result, diagnostics.render(
report = diagnostics.render(
params["cwb_preset"],
len(model_names),
params.get("cwb_config") is not None,
)
embedding_report = params.get("_embedding_coalesce_report")
if embedding_report is not None:
report += embedding_report.render()
return result, report
return result
finally:
for handler in handlers:
handler.__exit__(None, None, None)
if boundary_handler is not None:
boundary_handler.__exit__(None, None, None)
cleanup_after_operation()
@staticmethod
@@ -1486,6 +1671,7 @@ class ConsensusMergerLogic:
settings,
*,
embedding_union,
boundary_handler=None,
):
primary = handlers[0]
keys = set()
@@ -1579,6 +1765,41 @@ class ConsensusMergerLogic:
continue
raw = {i: loaded[(i, key)] for i in source_indices}
embedding_coalesce = bool(params.get("embedding_coalesce", False))
vision_boundaries = bool(params.get("vision_boundary_embeddings", False))
boundary_scores = []
trimmed_rows = 0
boundary_start = boundary_end = None
if embedding_coalesce and vision_boundaries:
if not all(tensor.ndim == 2 for tensor in raw.values()):
raise ValueError(
f"Vision boundary mode requires 2D embedding tensor '{key}'."
)
if params.get("legacy_boundary_search", False):
if boundary_handler is None:
raise ValueError(
"Legacy vision boundary search requires a boundary reference embedding."
)
boundary_start, boundary_end = _read_visual_boundary_reference(
boundary_handler, key
)
else:
anchor = raw[preserve_index]
boundary_start = anchor[0].detach().clone()
boundary_end = anchor[-1].detach().clone()
prepared_raw = {}
for source_index, tensor in raw.items():
body, scores, trimmed = _prepare_vision_embedding(
tensor,
start_reference=boundary_start,
end_reference=boundary_end,
legacy_search=params.get("legacy_boundary_search", False),
threshold=float(params["boundary_similarity_threshold"]),
)
prepared_raw[source_index] = body
boundary_scores.extend(scores)
trimmed_rows += trimmed
raw = prepared_raw
reference_index = 0
if embedding_union:
reference_source = max(
@@ -1657,6 +1878,31 @@ class ConsensusMergerLogic:
operation_label=key,
diagnostics=params.get("_cwb_diagnostics"),
)
if embedding_coalesce and merged.ndim == 2:
before_rows = merged.shape[0]
merged, pairs = coalesce_embedding_rows(
merged,
settings,
target_vector_count=int(params["target_vector_count"]),
similarity_threshold=float(params["coalesce_similarity_threshold"]),
position_window=float(params["coalesce_position_window"]),
)
if vision_boundaries:
merged = torch.cat((
boundary_start.to(dtype=merged.dtype, device=merged.device).unsqueeze(0),
merged,
boundary_end.to(dtype=merged.dtype, device=merged.device).unsqueeze(0),
), dim=0)
report = params.get("_embedding_coalesce_report")
if report is not None:
report.record(
key,
before_rows + (2 if vision_boundaries else 0),
merged.shape[0],
trimmed_rows,
boundary_scores,
pairs,
)
writer.write_batch([(key, merged)])
del merged
if secondary_only:
@@ -2517,6 +2763,118 @@ class CWBEmbeddingThreeMerger(CWBEmbeddingTwoMerger):
DEFAULT_FILENAME = "cwb_merged_3_embedding"
def _embedding_coalesce_controls(default_filename: str):
embedding_options = folder_paths.get_filename_list("embeddings")
return [
io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"], default="MERGE", tooltip="MERGE writes a new embedding. DOCUMENTATION ONLY opens no files."),
io.Combo.Input("cwb_preset", options=list(EMBEDDING_CWB_PRESETS), default="balanced_sim_mean", tooltip="Embedding CWB preset used for multi-input alignment and coalesced row pairs."),
CWB_CONFIG.Input("cwb_config", optional=True, tooltip="Optional complete CWB configuration override."),
io.Int.Input("target_vector_count", default=0, min=0, max=100000, step=1, tooltip="Target body-row count. Zero coalesces every eligible pair; otherwise stops at this count or when no eligible pair remains."),
io.Float.Input("coalesce_similarity_threshold", default=0.95, min=0.0, max=1.0, step=0.01, tooltip="Minimum cosine similarity for a mutual-nearest body-row pair to coalesce."),
io.Float.Input("coalesce_position_window", default=0.10, min=0.0, max=1.0, step=0.01, tooltip="Maximum normalized row-position distance for an eligible pair."),
io.Boolean.Input("vision_boundary_embeddings", default=False, tooltip="Enable for complete encoded visual blocks. Start/end rows are validated and preserved; False treats the full tensor as open textual inversion."),
io.Boolean.Input("legacy_boundary_search", default=False, tooltip="Search and trim legacy template rows using Boundary Reference Embedding. Requires vision boundary mode and a selected reference."),
io.Combo.Input("boundary_reference_embedding", options=["None", *embedding_options], default="None", tooltip="Known-clean visual block used by legacy boundary search. Its first and last encoded rows define vision start/end."),
io.Float.Input("boundary_similarity_threshold", default=0.95, min=-1.0, max=1.0, step=0.01, tooltip="Minimum encoded cosine similarity required independently for visual start and end validation."),
io.String.Input("output_filename", default=default_filename, tooltip="Filename without extension under ComfyUI's embeddings category."),
io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"], default="fp32", tooltip="Requested dtype for generated floating embedding tensors."),
io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda", tooltip="Per-tensor FP32 CWB arithmetic device; CUDA OOM retries that tensor on CPU."),
io.Boolean.Input("lazy_load", default=True, tooltip="Use bounded UEL streaming and release each tensor after its work unit."),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear Python and CUDA caches before each tensor at a possible speed cost."),
io.Boolean.Input("override_dtype", default=False, tooltip="Force generated floating tensors to Save Dtype."),
]
def _embedding_coalesce_params(kwargs: dict) -> dict:
return {
**kwargs,
"mismatch_mode": "error",
"exclude_patterns": "",
"discard_patterns": "",
"glob_patterns": False,
"include_1d_diffs": False,
"embedding_coalesce": True,
"_embedding_coalesce_report": EmbeddingCoalesceReport(
target_vector_count=int(kwargs["target_vector_count"]),
similarity_threshold=float(kwargs["coalesce_similarity_threshold"]),
position_window=float(kwargs["coalesce_position_window"]),
vision_boundary_embeddings=bool(kwargs["vision_boundary_embeddings"]),
legacy_boundary_search=bool(kwargs["legacy_boundary_search"]),
boundary_similarity_threshold=float(kwargs["boundary_similarity_threshold"]),
boundary_reference=None if kwargs["boundary_reference_embedding"] == "None" else kwargs["boundary_reference_embedding"],
),
"_return_cwb_report": True,
}
class CWBEmbeddingSelfCoalesce(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="CWBEmbeddingSelfCoalesce",
display_name="CWB Embedding Self-Coalesce",
category="ModelUtils/Embeddings/Merge",
description="Reduce one embedding's locally similar encoded vectors with CWB while preserving optional visual boundaries.",
inputs=[
io.Combo.Input("embedding", options=folder_paths.get_filename_list("embeddings"), tooltip="Embedding to coalesce."),
*_embedding_coalesce_controls("cwb_coalesced_embedding"),
],
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):
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.")
if kwargs["legacy_boundary_search"] and not kwargs["vision_boundary_embeddings"]:
raise ValueError("Legacy boundary search requires Vision Boundary Embeddings.")
filename, report = ConsensusMergerLogic.execute(
[kwargs["embedding"]], "embeddings", _embedding_coalesce_params(kwargs), embedding_union=True
)
return io.NodeOutput(filename, documentation, report)
class CWBEmbeddingMultiMerger(io.ComfyNode):
@classmethod
def define_schema(cls):
embeddings = folder_paths.get_filename_list("embeddings")
optional_embeddings = ["None", *embeddings]
return io.Schema(
node_id="CWBEmbeddingMultiMerger",
display_name="CWB Embedding Multi-Merge",
category="ModelUtils/Embeddings/Merge",
description="CWB-merge 2 to 8 embeddings, then coalesce locally similar output vectors.",
inputs=[
io.Combo.Input("embedding_count", options=[str(index) for index in range(2, 9)], default="2", tooltip="Number of consecutive embedding inputs to merge."),
io.Combo.Input("embedding_1", options=embeddings, tooltip="Primary embedding and normal visual-boundary anchor."),
io.Combo.Input("embedding_2", options=embeddings, tooltip="Second equal-prior embedding contributor."),
*[io.Combo.Input(f"embedding_{index}", options=optional_embeddings, default="None", tooltip="Optional embedding contributor.") for index in range(3, 9)],
*_embedding_coalesce_controls("cwb_merged_coalesced_embedding"),
],
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):
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.")
if kwargs["legacy_boundary_search"] and not kwargs["vision_boundary_embeddings"]:
raise ValueError("Legacy boundary search requires Vision Boundary Embeddings.")
count = int(kwargs["embedding_count"])
names = [kwargs[f"embedding_{index}"] for index in range(1, count + 1)]
missing = [str(index) for index, name in enumerate(names, start=1) if not name or name == "None"]
if missing:
raise ValueError(f"Embedding Count includes unselected input(s): {', '.join(missing)}.")
filename, report = ConsensusMergerLogic.execute(
names, "embeddings", _embedding_coalesce_params(kwargs), embedding_union=True
)
return io.NodeOutput(filename, documentation, report)
CWB_MERGER_NODES = [
CWBCustomConfiguration,
CWBCheckpointTwoMerger,
@@ -2530,4 +2888,6 @@ CWB_MERGER_NODES = [
CWBLoRAMultiMerger,
CWBEmbeddingTwoMerger,
CWBEmbeddingThreeMerger,
CWBEmbeddingSelfCoalesce,
CWBEmbeddingMultiMerger,
]
tests/test_consensus_merger.py
+133 -2
View File
@@ -588,9 +588,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) == 12
assert len(cwb.CWB_MERGER_NODES) == 14
schemas = [node.define_schema() for node in cwb.CWB_MERGER_NODES]
assert len({schema.node_id for schema in schemas}) == 12
assert len({schema.node_id for schema in schemas}) == 14
advanced = {
"consensus_type", "alignment_method", "alignment_threshold",
"similarity_threshold", "power_alpha", "diversity_beta",
@@ -623,6 +623,13 @@ def test_all_cwb_schemas_and_compact_control_contract(cwb):
"counterfactual_weight_sweep",
]
continue
if node in (cwb.CWBEmbeddingSelfCoalesce, cwb.CWBEmbeddingMultiMerger):
assert "vision_boundary_embeddings" in ids
assert "legacy_boundary_search" in ids
assert "boundary_reference_embedding" in ids
assert "boundary_similarity_threshold" in ids
assert "target_vector_count" in ids
continue
assert ids[:3] == ["execution_mode", "model_a", "model_b"]
if node.INPUT_COUNT == 3:
assert ids[3] == "model_c"
@@ -903,6 +910,130 @@ def test_lora_multi_wrapper_validates_count_and_forwards_equal_prior_names(
assert result.result == ("merged.safetensors", "DOCS", "REPORT")
def test_embedding_coalescing_uses_mutual_local_pairs(cwb):
rows = torch.tensor([
[1.0, 0.0],
[0.99, 0.01],
[0.0, 1.0],
])
merged, pairs = cwb.coalesce_embedding_rows(
rows,
_settings(cwb),
target_vector_count=2,
similarity_threshold=0.95,
position_window=1.0,
)
assert merged.shape == (2, 2)
assert pairs == ((0, 1),)
torch.testing.assert_close(merged[1], rows[2])
def test_visual_boundary_preparation_validates_and_searches_legacy_rows(cwb):
start = torch.tensor([1.0, 0.0])
end = torch.tensor([0.0, 1.0])
body = torch.tensor([[0.6, 0.4]])
legacy = torch.stack((torch.tensor([-1.0, 0.0]), start, body[0], end, torch.tensor([0.0, -1.0])))
prepared, scores, trimmed = cwb._prepare_vision_embedding(
legacy,
start_reference=start,
end_reference=end,
legacy_search=True,
threshold=0.99,
)
torch.testing.assert_close(prepared, body)
assert scores == pytest.approx((1.0, 1.0))
assert trimmed == 2
with pytest.raises(ValueError, match="Vision boundary validation failed"):
cwb._prepare_vision_embedding(
torch.stack((start, body[0], -end)),
start_reference=start,
end_reference=end,
legacy_search=False,
threshold=0.99,
)
def test_embedding_self_coalesce_streams_and_restores_visual_boundaries(
monkeypatch, tmp_path, cwb
):
source = tmp_path / "source.safetensors"
start = torch.tensor([1.0, 0.0])
end = torch.tensor([0.0, 1.0])
save_file({"visual": torch.stack((start, torch.tensor([0.99, 0.01]), torch.tensor([0.98, 0.02]), end))}, str(source))
_patch_io(monkeypatch, cwb, tmp_path, {"source": str(source)})
monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS")
result = cwb.CWBEmbeddingSelfCoalesce.execute(
embedding="source",
execution_mode="MERGE",
cwb_preset="balanced_sim_mean",
cwb_config=_settings(cwb),
target_vector_count=1,
coalesce_similarity_threshold=0.95,
coalesce_position_window=1.0,
vision_boundary_embeddings=True,
legacy_boundary_search=False,
boundary_reference_embedding="None",
boundary_similarity_threshold=0.99,
output_filename="visual_coalesced",
save_dtype="fp32",
process_device="cpu",
lazy_load=True,
force_clear_cache=False,
override_dtype=False,
)
output = load_file(str(_result_path(tmp_path, "embeddings", result.result[0])))
assert output["visual"].shape == (3, 2)
torch.testing.assert_close(output["visual"][0], start)
torch.testing.assert_close(output["visual"][-1], end)
assert "CWB EMBEDDING COALESCING" in result.result[2]
assert "Similarity threshold: 0.9500" in result.result[2]
def test_embedding_multi_wrapper_forwards_selected_inputs(monkeypatch, cwb):
kwargs = {
"execution_mode": "MERGE",
"embedding_count": "3",
"embedding_1": "a.safetensors",
"embedding_2": "b.safetensors",
"embedding_3": "c.safetensors",
"embedding_4": "None",
"embedding_5": "None",
"embedding_6": "None",
"embedding_7": "None",
"embedding_8": "None",
"cwb_preset": "balanced_sim_mean",
"cwb_config": None,
"target_vector_count": 0,
"coalesce_similarity_threshold": 0.95,
"coalesce_position_window": 0.1,
"vision_boundary_embeddings": False,
"legacy_boundary_search": False,
"boundary_reference_embedding": "None",
"boundary_similarity_threshold": 0.95,
"output_filename": "merged",
"save_dtype": "fp32",
"process_device": "cpu",
"lazy_load": True,
"force_clear_cache": False,
"override_dtype": False,
}
monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS")
captured = {}
def execute(names, model_type, params, **options):
captured.update(names=names, model_type=model_type, params=params, options=options)
return "merged.safetensors", "REPORT"
monkeypatch.setattr(cwb.ConsensusMergerLogic, "execute", execute)
result = cwb.CWBEmbeddingMultiMerger.execute(**kwargs)
assert captured["names"] == ["a.safetensors", "b.safetensors", "c.safetensors"]
assert captured["model_type"] == "embeddings"
assert captured["options"] == {"embedding_union": True}
assert captured["params"]["embedding_coalesce"] is True
assert result.result == ("merged.safetensors", "DOCS", "REPORT")
def test_streaming_dtype_nonfloat_and_secondary_preservation(
monkeypatch, tmp_path, cwb
):