diff --git a/audit/uel_multi_input_loading_audit.md b/audit/uel_multi_input_loading_audit.md deleted file mode 100644 index c68b07f..0000000 --- a/audit/uel_multi_input_loading_audit.md +++ /dev/null @@ -1,136 +0,0 @@ -# Initial UEL Loading Audit - -Multi-input loading is inconsistent across the repository. Model analysis and -LoRA multi-merge use bounded UEL asynchronous streams with pinned-memory -support, while LoRA extraction, generic model merging, CWB merging, and parts -of LoRA-to-model application still call `get_tensor()` sequentially for each -input. - -LoRA Knee Detection is especially affected because synchronous A/B reads are -combined with an exact full FP32 SVD per eligible layer. Its `max_rank` limit is -applied only after the full decomposition. The extractor also builds a full -residual reconstruction that its caller discards, clears caches after every -layer by default, and attempts fused-layer chunking only after the initial SVD -fails. - -Initial recommendation: audit all multi-input operations against the bounded -UEL work-unit pattern already used by model analysis and LoRA multi-merge, then -migrate the shared execution paths rather than patching Knee Detection alone. - -## Knee SVD follow-up - -Knee extraction now uses a configurable partial-spectrum probe beginning at -`max_rank + knee_probe_offset`. A knee detected in the probe tail triggers one -expansion up to `2 * max_rank`; saved rank remains capped by `max_rank`. This is -shared by LoRA, DoRA, learned DoRA, and both text-encoder knee node families. -The same change removes discarded full residual reconstruction and replaces -diagonal-matrix multiplication with direct singular-value column scaling. - -## Recovery contract - -This is a working-tree checkpoint, not proof that the implementation is -correct. On recovery, verify each item against the named code path and rerun -its focused test before relying on it. Passing tests prove only their -assertions. Do not infer lifecycle, release behavior, alpha mathematics, or -repository-wide coverage from an import, helper name, or checkbox. - -## Repository-wide migration progress - -- [x] Audit production `get_tensor()` and direct safetensors call sites. -- [ ] Verify the shared bounded UEL stream and atomic writer implementation. -- [ ] Verify generic and CWB merging has no hidden synchronous tensor I/O. -- [ ] Verify resize-with-base and resize use ordered UEL work units. -- [ ] Verify key utilities, metadata, passthrough, and PT conversion output. -- [ ] Verify the repository guard covers every production tensor-I/O path. -- [ ] Verify alpha normalization precedes every LoRA rank operation. -- [ ] Finish atomic-writer migration and remove obsolete imports/helpers. -- [ ] Add repository guards and complete lifecycle/failure tests. -- [ ] Run focused tests, full validation, and representative benchmarks. - -## Unverified working-tree changes - -The working tree intends `nodes/uel_io.py` to provide ordered streams, logical -multi-input work units, per-yield release, stream closure, and atomic output. -Generic recipes were changed to receive preloaded mappings. CWB, LoRA-to-model, -key utilities, metadata, and PT conversion were changed toward UEL. These are -implementation claims awaiting focused lifecycle verification. - -The intended alpha contract is `up * (alpha / rank)` before padding, alignment, -weighting, CWB, DARE/TIES, or model application. Callers were changed to replace -the original work-unit reference and omit alpha outputs. Verify every caller, -including preservation and early-return branches. The standalone normalization -script should remain untracked and unchanged; confirm this before staging. - -Evidence recorded so far: focused CWB and alpha tests passed. The first full -run reached 152 passing tests -with two failures caused only by tests reaching an implementation module's -removed `IncrementalSafetensorsWriter` import. Update those tests to import the -UEL writer directly, then rerun the full suite. A bounded search found no -production `get_tensor()` or direct safetensors imports at that point; rerun the -guard after every later edit. - -## Remaining execution checklist - -1. Model analysis was rechecked after an incorrect intermediate concern. The - caller near line 910 passes exactly one A/B tensor pair. The LoRA caller near - line 953 passes exactly two pairs for one down/up logical group, which must - coexist for pair analysis. It does not accumulate the model inventory. Both - callers release and mark their bounded unit before advancing. No change is - required unless a later test contradicts this inspection. -2. Re-run the new failure lifecycle test after explicit caller-side stream - closure changes in generic merge, CWB, and LoRA-to-model. -3. Verify extraction and resize atomic-writer conversions and remove obsolete - imports. Dtype conversion was switched to the shared atomic helper. -4. The quantization test fixture now imports the UEL writer directly. The next - full run reached 157 passing tests and one newly added lifecycle test failure; - that failure exposed missing explicit consumer-side generator closure and - code was changed afterward. Rerun the full suite. -5. The new repository guard and atomic failure test are in - `tests/test_uel_repository_contract.py`; verify they still pass. -6. Run Ruff F checks, compilation, `git diff --check`, full pytest, then a - temporary synthetic benchmark. Remove only task-created temp directories. -7. Reinspect every `async_stream()` caller for accumulation, release timing, - exceptional closure, and duplicate key occurrences. Imports and method names - are not evidence of a correct implementation. - -## Latest verification evidence - -- Full configured-environment suite: 158 tests passed. -- Focused alpha and UEL contract rerun after the full suite: 7 tests passed. -- Changed Python files compile with the configured ComfyUI interpreter. -- Ruff F checks pass for all changed implementation and test files. Unrelated - pre-existing F findings remain in untouched modules outside this change. -- `git diff --check` passes. -- The production guard finds no `.get_tensor()` calls or direct safetensors - imports under `nodes/*.py`. -- No em dash occurs in the changed audit, implementation, or focused test files. - -These results are evidence for this exact working-tree revision only. Rerun -them after any edit, merge, formatting pass, or conflict resolution. -## Compaction checkpoint: alpha normalization edge cases - -Verified on 2026-08-14 by bounded source inspection: - -- `nodes/merger.py` now records `(alpha_key, rank, down_key)` for each alpha-bearing LoRA pair and rejects alpha normalization when either factor is an unsupported low-bit tensor. Its runtime scales the loaded up factor before merge processing. -- The remaining work is to add the same explicit rejection to LoRA multi-merge and CWB LoRA guarded low-bit paths. Scalar low-bit alpha alone must remain loadable; only unsupported low-bit factor tensors make mathematical normalization impossible. -- `write_preserved_tensor` still needs failure-path lifecycle verification so every yielded tensor is marked processed and its stream is closed if writer submission raises. -- After those edits, rerun focused low-bit/alpha tests, the complete test suite in the configured ComfyUI virtual environment, changed-file compile and Ruff checks, `git diff --check`, and the production guard against direct `get_tensor()` or direct safetensors imports. -- Do not stage this work unless explicitly requested. Keep all unrelated untracked files listed by `git status` untracked. - -## Fresh validation evidence after recovery - -The following is command evidence from 2026-08-14, not a substitute for inspecting the current tree after later edits: - -- The configured ComfyUI Python completed the full suite with `160 passed in 5.52s`. -- Focused regressions for CWB and LoRA multi-merge rejection of alpha-bearing unsupported low-bit factors both passed. -- Changed implementation and test files passed Ruff `F` checks through the installed standalone Ruff executable. -- `python -m compileall -q nodes tests` completed successfully. -- `git diff --check` completed with no diff errors; Git emitted only line-ending conversion warnings. -- A scoped production search found no direct `.get_tensor(` calls and no direct safetensors imports in `nodes/*.py`. -- A scoped search found no em dash in the UEL audit, shared helpers, affected merger files, or UEL contract test. - -Implementation consequences verified by the focused source inspection and tests above: - -- Alpha-bearing LoRA factors are normalized before merge and before padding when their factor tensors are supported floating storage. -- A low-bit scalar alpha is accepted when both factors are supported; unsupported low-bit down/up factors with alpha are rejected instead of being copied into an output that falsely claims normalization. -- Preserved tensors yielded by the one-key UEL stream are marked processed in `finally`, including writer-failure paths, and the stream is closed. diff --git a/docs/consensus_mergers.md b/docs/consensus_mergers.md index 87d26b2..b3a78d7 100644 --- a/docs/consensus_mergers.md +++ b/docs/consensus_mergers.md @@ -14,7 +14,8 @@ conditioning and spatial fusion, which are not controls of these nodes. ## Node families Two-input and three-input variants are provided for checkpoints, standalone -diffusion models, text encoders, LoRAs, and embeddings. +diffusion models, text encoders, LoRAs, and embeddings. LoRA also provides a +Multi-Merge variant for 2 to 8 equal-prior inputs. 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 @@ -26,12 +27,18 @@ 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. Every named merger preset keeps `global_scale` at `1.0`; scaling -persisted model or LoRA weights is only performed when explicitly selected in -the custom controls. +Presets are separated by merge type. Dense tensors, embeddings, and LoRA +factor pairs never share one menu. Every named preset supplies every CWB +setting and cannot inherit hidden widget state. + +Manual controls live in the separate CWB Custom Configuration node. Connecting +its `CWB_CONFIG` output completely overrides the selected preset. Leaving it +unconnected uses the named preset. + +Preset names expose categorical and boolean behavior. `_idx` and `_sim` +identify alignment, `_medn` and `_mean` identify consensus, and `_rn`, `_dsc`, +`_softcb`, and `_pcp` identify enabled boolean transformations. Dense presets +omit alignment tokens because dense model coordinates are fixed. Similarity alignment greedily matches source rows to the reference by cosine similarity. Index alignment keeps absolute row positions. Position weight can @@ -68,13 +75,8 @@ For each aligned vector group, the merger performs these operations in order: - **Model B / Model C**: additional equal-prior contributors. CWB derives their effective per-vector influence from consensus similarity; these nodes do not expose model-level weights. -- **CWB Preset**: `custom` makes all manual CWB controls authoritative. A named - preset sets consensus, similarity cutoff, alpha/beta, norm rescaling, and its - DSC fields. Named presets force similarity alignment. Only `power_blend` - supplies an alignment threshold (0.9); other presets retain the manual - alignment threshold. `position_weight` and `preserve_common_prefix` remain - explicit extensions. A manually selected `global_scale` other than 1.0 also - remains active. +- **CWB Preset**: complete use-case settings selected from the node's own dense, + embedding, or LoRA registry. A connected CWB Config overrides it completely. - **Consensus Type**: `mean` uses the arithmetic center; `median` uses the element-wise median and is less sensitive to coordinate outliers. - **Alignment Method**: `index` pairs first-axis vectors by position; @@ -154,26 +156,44 @@ in tensors with at least two dimensions, not to those scalar/1D paths. ### Presets implemented by these nodes -All named file-merger presets use a neutral `global_scale` of 1.0. This is an -intentional repository safety difference from the broad specification's 0.7 -calibration for some conditioning-oriented presets: persisted model and LoRA -strength is not reduced unless the user explicitly changes Global Scale. +Dense nodes offer `balanced_mean`, `robust_medn`, `selective_mean`, +`varied_mean_rn_softcb`, `diverse_medn_rn_dsc_softcb`, and +`strongdiv_medn_rn_dsc_softcb`. -| Preset | Consensus | Alpha | Similarity cutoff | Beta | Norm rescale | DSC | Soft comfort | Alignment cutoff | -| --- | --- | ---: | ---: | ---: | --- | --- | --- | --- | -| `baseline` | median | 2.0 | 0.0 | 0.0 | no | no | no | manual value | -| `power_blend` | median | 8.0 | 0.75 | 0.0 | yes | yes | no | 0.9 | -| `high_clarity` | median | 3.0 | 0.3 | 0.0 | no | no | no | manual value | -| `smooth` | mean | 1.5 | 0.0 | 0.0 | no | no | no | manual value | -| `varied_merge` | median | 2.0 | 0.0 | 0.0 | yes | no | no | manual value | -| `diverse_concept` | median | 2.0 | 0.0 | 1.0 | yes | no | no | manual value | -| `high_diversity_concept` | median | 2.0 | 0.0 | 2.0 | yes | no | no | manual value | -| `dsc_baseline` | median | 2.0 | 0.0 | 0.0 | no | yes | yes | manual value | -| `dsc_high_clarity` | median | 4.0 | 0.3 | 0.0 | no | yes | yes | manual value | -| `dsc_smooth` | mean | 1.0 | 0.0 | 0.0 | no | yes | yes | manual value | -| `dsc_varied_merge` | median | 2.5 | 0.0 | 0.0 | yes | yes | yes | manual value | -| `dsc_diverse_concept` | median | 2.0 | 0.0 | 1.5 | yes | yes | yes | manual value | -| `dsc_high_diversity_concept` | median | 2.0 | 0.0 | 3.0 | yes | yes | yes | manual value | +Embedding nodes offer `balanced_idx_mean`, `balanced_sim_mean`, +`robust_idx_medn`, `robust_sim_medn`, `varied_sim_mean_rn_softcb`, and +`diverse_sim_medn_rn_dsc_softcb`. + +LoRA nodes offer `broad_sim_medn_rn_softcb`, +`moderate_sim_medn_rn_softcb`, `conservative_sim_medn_rn_softcb`, +`direct_idx_medn_rn_softcb`, `broad_sim_mean_rn_softcb`, +`neutral_sim_medn_rn`, `focused_sim_medn_rn_dsc_softcb`, and +`strongfocus_sim_medn_rn_dsc_softcb`. + +The LoRA default is `broad_sim_medn_rn_softcb`: median consensus, +similarity alignment, zero similarity threshold, alpha 2, beta 4, norm +rescaling, scale 1, DSC disabled, soft comfort bandpass enabled, position +weight 0.05, and common-prefix preservation disabled. Broad, moderate, and +conservative use alignment thresholds 0, 0.0005, and 0.0025. They represent +increasing anchor preservation in automated mixed-rank LoRA tests. + +Every LoRA preset fixes all 12 controls explicitly: + +| preset | consensus | alignment | align threshold | sim threshold | alpha | beta | RN | scale | DSC | soft CB | position | prefix | intent | +|---|---|---|---:|---:|---:|---:|---|---:|---|---|---:|---|---| +| `broad_sim_medn_rn_softcb` | median | similarity | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Broad row matching; default | +| `moderate_sim_medn_rn_softcb` | median | similarity | 0.0005 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Moderate anchor preservation | +| `conservative_sim_medn_rn_softcb` | median | similarity | 0.0025 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Strong anchor preservation | +| `direct_idx_medn_rn_softcb` | median | index | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Fixed rank-row correspondence | +| `broad_sim_mean_rn_softcb` | mean | similarity | 0 | 0 | 2 | 4 | on | 1 | off | on | 0.05 | off | Mean-consensus broad merge | +| `neutral_sim_medn_rn` | median | similarity | 0 | 0 | 2 | 0 | on | 1 | off | off | 0.05 | off | Similarity weighting without diversity modulation | +| `focused_sim_medn_rn_dsc_softcb` | median | similarity | 0 | 0 | 2 | 4 | on | 1 | on | on | 0.05 | off | Moderate DSC concentration | +| `strongfocus_sim_medn_rn_dsc_softcb` | median | similarity | 0 | 0 | 2 | 7 | on | 1 | on | on | 0.05 | off | Strong DSC concentration | + +The LoRA Multi-Merge counterfactual weight sweep reuses each consensus vector +already computed by the active merge. It reports mean and median consensus +results across configured alpha, beta, similarity-threshold, DSC, and comfort +bandpass values without additional model loads or output files. ## Missing tensors and filtering @@ -200,15 +220,28 @@ direct `.diff`, `.diff_b`, `.w_norm`, `.b_norm`, and `.set_weight` layers. Recognized diffusion-model outputs use the canonical `diffusion_model..lora_A.weight` / `.lora_B.weight` convention. -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 anchor-source alpha key is updated -to the resulting maximum rank; no alpha key is invented when the anchor has -none. +Secondary formats are matched by normalized logical layer name. Alpha is +absorbed mathematically into each B/up factor using its original rank before +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. Similarity alignment operates on paired latent-rank components. A rows and B columns share one mapping; fixed input/output feature axes are never reordered. -Input alpha scaling is absorbed into the B factor before blending, and global -scale is applied once to the resulting pair rather than once per factor. +For a rank-R reference and rank-r source it evaluates an R-by-r similarity +matrix, reduced only by a preserved common prefix. Every genuine source +component can therefore match any genuine reference component. Index alignment +maps source component i to reference component i. In either mode, reference +components without a genuine contributor remain singleton groups and are not +weighted or norm-rescaled against nonexistent rank padding. + +Structural rank absence is different from `mismatch_mode=zeros`. The latter is +an explicit request for a semantic zero contributor and continues to apply to +every output-rank component where the mismatch policy is active. Numerical zero +values inside a genuine factor do not identify padding; validity comes from the +factor's original rank. + +Global scale is applied once to the resulting pair rather than once per factor. Incomplete pairs and unrecognized tensors are copied from the earliest available source unchanged. Companion-bearing groups (`lora_mid`, reshape, DoRA, or set-weight) are kept @@ -235,6 +268,40 @@ 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. +## CWB report output + +Every CWB merge node provides a separate `cwb_report` string output after the +existing filename and documentation outputs. Documentation-only execution +returns an explicit message that no report exists because no tensors were +processed. + +The filename output uses ComfyUI's wildcard socket type and returns the path +relative to the matching ComfyUI model category. A filename such as +`groupfolder/modelname` therefore returns +`groupfolder/modelname.safetensors`. It can connect directly to compatible +loader Combo inputs without exposing the absolute models-directory path. + +The report begins with model-wide counters. LoRA arithmetic then adds grouped +mathematical evidence without retaining tensors or dumping complete similarity +matrices. Layers are grouped only when they share: + +- down non-rank dimensions; +- up non-rank dimensions; +- ordered genuine input ranks; +- selected reference input and rank; +- alignment method. + +Each dimension group reports layer count, effective matrix dimensions, genuine +source-component match coverage, reference-only components, structural rank +slots excluded from arithmetic, explicit mismatch-zero contributors, candidate +similarity min/mean/max, accepted-match similarity percentiles, and rank-one +delta norm ratios. A compact exception section identifies the lowest match +coverage and largest norm changes by layer name. + +For a rank-384 reference and rank-128 source, the report identifies a `384x128` +matrix, at most 128 genuine matches, and the 256 structural rank slots that were +excluded rather than treated as zero-valued contributors. + ## Streaming and failure safety Model parameters and direct patches use fixed-coordinate blending; similarity diff --git a/nodes/consensus_merger.py b/nodes/consensus_merger.py index c87fed9..03fb493 100644 --- a/nodes/consensus_merger.py +++ b/nodes/consensus_merger.py @@ -5,7 +5,7 @@ from __future__ import annotations import logging import os from contextlib import closing -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Iterable import comfy.utils @@ -22,6 +22,7 @@ from .device_utils import ( prepare_for_large_operation, ) from .uel_io import atomic_uel_writer, stream_work_units +from .artifact_paths import canonical_model_artifact_path from .lora_alpha import normalize_lora_pair from .lora_resize import ( canonical_lora_key, @@ -42,83 +43,195 @@ from .quantization_guard import ( ) -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": 1.0, - }, - "diverse_concept": { - "consensus_type": "median", "power_alpha": 2.0, - "similarity_threshold": 0.0, "diversity_beta": 1.0, - "rescale_norm": True, "global_scale": 1.0, - }, - "high_diversity_concept": { - "consensus_type": "median", "power_alpha": 2.0, - "similarity_threshold": 0.0, "diversity_beta": 2.0, - "rescale_norm": True, "global_scale": 1.0, - }, - "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": 1.0, - "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": 1.0, - "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": 1.0, - "dynamic_similarity_contrast": True, "soft_comfort_bandpass": True, - }, +def _preset( + *, consensus_type, alignment_method, alignment_threshold, + similarity_threshold, power_alpha, diversity_beta, + rescale_norm, global_scale, dynamic_similarity_contrast, + soft_comfort_bandpass, position_weight, preserve_common_prefix, +): + return { + "consensus_type": consensus_type, + "alignment_method": alignment_method, + "alignment_threshold": alignment_threshold, + "similarity_threshold": similarity_threshold, + "power_alpha": power_alpha, + "diversity_beta": diversity_beta, + "rescale_norm": rescale_norm, + "global_scale": global_scale, + "dynamic_similarity_contrast": dynamic_similarity_contrast, + "soft_comfort_bandpass": soft_comfort_bandpass, + "position_weight": position_weight, + "preserve_common_prefix": preserve_common_prefix, + } + + +DENSE_CWB_PRESETS = { + "balanced_mean": _preset( + consensus_type="mean", alignment_method="index", alignment_threshold=0.0, + similarity_threshold=0.0, power_alpha=1.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, + ), + "robust_medn": _preset( + consensus_type="median", alignment_method="index", alignment_threshold=0.0, + 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, + ), + "selective_mean": _preset( + consensus_type="mean", alignment_method="index", alignment_threshold=0.0, + similarity_threshold=0.3, power_alpha=3.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, + ), + "varied_mean_rn_softcb": _preset( + consensus_type="mean", alignment_method="index", alignment_threshold=0.0, + similarity_threshold=0.0, power_alpha=1.5, diversity_beta=2.0, + rescale_norm=True, global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=True, position_weight=0.0, + preserve_common_prefix=False, + ), + "diverse_medn_rn_dsc_softcb": _preset( + consensus_type="median", alignment_method="index", alignment_threshold=0.0, + similarity_threshold=0.0, power_alpha=2.0, diversity_beta=4.0, + rescale_norm=True, global_scale=1.0, dynamic_similarity_contrast=True, + soft_comfort_bandpass=True, position_weight=0.0, + preserve_common_prefix=False, + ), + "strongdiv_medn_rn_dsc_softcb": _preset( + consensus_type="median", alignment_method="index", alignment_threshold=0.0, + similarity_threshold=0.0, power_alpha=2.0, diversity_beta=10.0, + rescale_norm=True, global_scale=1.0, dynamic_similarity_contrast=True, + soft_comfort_bandpass=True, position_weight=0.0, + preserve_common_prefix=False, + ), } -CWB_PRESET_OPTIONS = ["custom", *CWB_PRESETS] +EMBEDDING_CWB_PRESETS = { + "balanced_idx_mean": _preset( + consensus_type="mean", alignment_method="index", alignment_threshold=0.0, + similarity_threshold=0.0, power_alpha=1.5, 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, + ), + "balanced_sim_mean": _preset( + consensus_type="mean", alignment_method="similarity", alignment_threshold=0.4, + similarity_threshold=0.0, power_alpha=1.5, diversity_beta=0.0, + rescale_norm=False, global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=False, position_weight=0.1, + preserve_common_prefix=False, + ), + "robust_idx_medn": _preset( + consensus_type="median", alignment_method="index", alignment_threshold=0.0, + 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, + ), + "robust_sim_medn": _preset( + consensus_type="median", alignment_method="similarity", + 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.1, + preserve_common_prefix=False, + ), + "varied_sim_mean_rn_softcb": _preset( + consensus_type="mean", alignment_method="similarity", + alignment_threshold=0.2, similarity_threshold=0.0, + power_alpha=1.5, diversity_beta=2.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=True, position_weight=0.1, + preserve_common_prefix=False, + ), + "diverse_sim_medn_rn_dsc_softcb": _preset( + consensus_type="median", alignment_method="similarity", + alignment_threshold=0.0, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=4.0, rescale_norm=True, + global_scale=1.0, + dynamic_similarity_contrast=True, soft_comfort_bandpass=True, + position_weight=0.1, preserve_common_prefix=False, + ), +} + +LORA_CWB_PRESETS = { + "broad_sim_medn_rn_softcb": _preset( + consensus_type="median", alignment_method="similarity", + alignment_threshold=0.0, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=4.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=True, position_weight=0.05, + preserve_common_prefix=False, + ), + "moderate_sim_medn_rn_softcb": _preset( + consensus_type="median", alignment_method="similarity", + alignment_threshold=0.0005, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=4.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=True, position_weight=0.05, + preserve_common_prefix=False, + ), + "conservative_sim_medn_rn_softcb": _preset( + consensus_type="median", alignment_method="similarity", + alignment_threshold=0.0025, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=4.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=True, position_weight=0.05, + preserve_common_prefix=False, + ), + "direct_idx_medn_rn_softcb": _preset( + consensus_type="median", alignment_method="index", + alignment_threshold=0.0, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=4.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=True, position_weight=0.05, + preserve_common_prefix=False, + ), + "broad_sim_mean_rn_softcb": _preset( + consensus_type="mean", alignment_method="similarity", + alignment_threshold=0.0, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=4.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=True, position_weight=0.05, + preserve_common_prefix=False, + ), + "neutral_sim_medn_rn": _preset( + consensus_type="median", alignment_method="similarity", + alignment_threshold=0.0, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=0.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=False, + soft_comfort_bandpass=False, position_weight=0.05, + preserve_common_prefix=False, + ), + "focused_sim_medn_rn_dsc_softcb": _preset( + consensus_type="median", alignment_method="similarity", + alignment_threshold=0.0, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=4.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=True, + soft_comfort_bandpass=True, position_weight=0.05, + preserve_common_prefix=False, + ), + "strongfocus_sim_medn_rn_dsc_softcb": _preset( + consensus_type="median", alignment_method="similarity", + alignment_threshold=0.0, similarity_threshold=0.0, + power_alpha=2.0, diversity_beta=7.0, rescale_norm=True, + global_scale=1.0, dynamic_similarity_contrast=True, + soft_comfort_bandpass=True, position_weight=0.05, + preserve_common_prefix=False, + ), +} + + +def _presets_for(*, embedding_union: bool, lora_mode: bool): + if lora_mode: + return LORA_CWB_PRESETS + if embedding_union: + return EMBEDDING_CWB_PRESETS + return DENSE_CWB_PRESETS LORA_PREFIXES = ( "base_model.model.", "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", @@ -142,44 +255,508 @@ class CWBSettings: 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)), - } +def _optional_min(current: float | None, value: float | None) -> float | None: + if value is None: + return current + return value if current is None else min(current, value) + + +def _optional_max(current: float | None, value: float | None) -> float | None: + if value is None: + return current + return value if current is None else max(current, value) + + +COUNTERFACTUAL_SIMILARITY_THRESHOLDS = (0.0, 0.1, 0.25) +COUNTERFACTUAL_POWER_ALPHAS = (1.0, 2.0, 4.0) +COUNTERFACTUAL_DIVERSITY_BETAS = (0.0, 2.0, 4.0, 7.0, 10.0) + + +@dataclass +class CWBWeightSweepStat: + groups: int = 0 + contributors: int = 0 + accepted: int = 0 + all_rejected: int = 0 + zero_weight: int = 0 + dominant_sum: float = 0.0 + dominant_max: float = 0.0 + effective_sum: float = 0.0 + effective_min: float | None = None + + def absorb(self, other: "CWBWeightSweepStat") -> None: + self.groups += other.groups + self.contributors += other.contributors + self.accepted += other.accepted + self.all_rejected += other.all_rejected + self.zero_weight += other.zero_weight + self.dominant_sum += other.dominant_sum + self.dominant_max = max(self.dominant_max, other.dominant_max) + self.effective_sum += other.effective_sum + self.effective_min = _optional_min(self.effective_min, other.effective_min) + + +@dataclass +class CWBWeightSweepDiagnostics: + stats: dict[tuple, CWBWeightSweepStat] = field(default_factory=dict) + + def record(self, stacked: torch.Tensor) -> None: + for consensus_type in ("mean", "median"): + consensus = ( + torch.mean(stacked, dim=0) + if consensus_type == "mean" + else torch.median(stacked, dim=0).values + ) + similarities = torch.mv( + F.normalize(stacked, p=2, dim=1, eps=1e-8), + F.normalize(consensus, p=2, dim=0, eps=1e-8), + ) + self._record_similarities(consensus_type, similarities) + + def _record_similarities( + self, + consensus_type: str, + similarities: torch.Tensor, + ) -> None: + values = similarities.detach().float().cpu().tolist() + minimum = min(values) + maximum = max(values) + for dsc in (False, True): + if dsc and maximum > minimum: + weighted = [ + 0.7 + 0.3 * (value - minimum) / (maximum - minimum + 1e-8) + for value in values + ] + else: + weighted = values + for threshold in COUNTERFACTUAL_SIMILARITY_THRESHOLDS: + accepted = [value >= threshold for value in values] + for alpha in COUNTERFACTUAL_POWER_ALPHAS: + for beta in COUNTERFACTUAL_DIVERSITY_BETAS: + for softcb in (False, True): + key = ( + consensus_type, threshold, alpha, beta, dsc, softcb + ) + stat = self.stats.setdefault(key, CWBWeightSweepStat()) + stat.groups += 1 + stat.contributors += len(values) + stat.accepted += sum(accepted) + if not any(accepted): + stat.all_rejected += 1 + weights = [1.0 / len(values)] * len(values) + else: + distance_base = 1.5 if softcb else 1.001 + weights = [] + for value, include in zip(weighted, accepted): + if not include: + weights.append(0.0) + continue + safe = min(max(value, 0.0), 1.0) + weight = safe ** alpha + if beta > 0.0: + weight *= max(distance_base - safe, 0.0) ** beta + weights.append(weight) + total = sum(weights) + if total <= 0.0: + stat.zero_weight += 1 + weights = [1.0 / len(values)] * len(values) + else: + weights = [weight / total for weight in weights] + dominant = max(weights) + effective = 1.0 / sum(weight * weight for weight in weights) + stat.dominant_sum += dominant + stat.dominant_max = max(stat.dominant_max, dominant) + stat.effective_sum += effective + stat.effective_min = _optional_min( + stat.effective_min, effective + ) + + def absorb(self, other: "CWBWeightSweepDiagnostics") -> None: + for key, other_stat in other.stats.items(): + self.stats.setdefault(key, CWBWeightSweepStat()).absorb(other_stat) + + def render(self) -> str: + lines = [ + "", + "CWB COUNTERFACTUAL WEIGHT SWEEP", + "consensus | sim_threshold | alpha | beta | dsc | softcb | accepted | fallbacks | dominant_mean/max | effective_mean/min", + ] + for key in sorted(self.stats): + consensus_type, threshold, alpha, beta, dsc, softcb = key + stat = self.stats[key] + lines.append( + f"{consensus_type} | {threshold:.6g} | {alpha:.6g} | {beta:.6g} | " + f"{str(dsc).lower()} | {str(softcb).lower()} | " + f"{stat.accepted}/{stat.contributors} | " + f"{stat.all_rejected + stat.zero_weight} | " + f"{stat.dominant_sum / stat.groups:.6g}/{stat.dominant_max:.6g} | " + f"{stat.effective_sum / stat.groups:.6g}/{stat.effective_min:.6g}" + ) + return "\n".join(lines) + + +@dataclass +class CWBDiagnostics: + alignment_candidates: int = 0 + alignment_matches: int = 0 + anchor_only_groups: int = 0 + weighting_groups: int = 0 + all_rejected_fallbacks: int = 0 + zero_weight_fallbacks: int = 0 + weighting_contributors: int = 0 + accepted_weighting_contributors: int = 0 + consensus_similarity_sum: float = 0.0 + consensus_similarity_min: float | None = None + consensus_similarity_max: float | None = None + dominant_weight_sum: float = 0.0 + dominant_weight_min: float | None = None + dominant_weight_max: float | None = None + effective_contributor_sum: float = 0.0 + effective_contributor_min: float | None = None + effective_contributor_max: float | None = None + cpu_fallbacks: int = 0 + lora_groups: int = 0 + mixed_rank_lora_groups: int = 0 + genuine_lora_components: int = 0 + structural_rank_slots_excluded: int = 0 + explicit_zero_contributors: int = 0 + lora_layer_reports: list["LoRALayerReport"] = field(default_factory=list) + weight_sweep: CWBWeightSweepDiagnostics | None = None + + def absorb_runtime(self, other: "CWBDiagnostics") -> None: + self.alignment_candidates += other.alignment_candidates + self.alignment_matches += other.alignment_matches + self.anchor_only_groups += other.anchor_only_groups + self.weighting_groups += other.weighting_groups + self.all_rejected_fallbacks += other.all_rejected_fallbacks + self.zero_weight_fallbacks += other.zero_weight_fallbacks + self.weighting_contributors += other.weighting_contributors + self.accepted_weighting_contributors += other.accepted_weighting_contributors + self.consensus_similarity_sum += other.consensus_similarity_sum + self.consensus_similarity_min = _optional_min( + self.consensus_similarity_min, other.consensus_similarity_min + ) + self.consensus_similarity_max = _optional_max( + self.consensus_similarity_max, other.consensus_similarity_max + ) + self.dominant_weight_sum += other.dominant_weight_sum + self.dominant_weight_min = _optional_min( + self.dominant_weight_min, other.dominant_weight_min + ) + self.dominant_weight_max = _optional_max( + self.dominant_weight_max, other.dominant_weight_max + ) + self.effective_contributor_sum += other.effective_contributor_sum + self.effective_contributor_min = _optional_min( + self.effective_contributor_min, other.effective_contributor_min + ) + self.effective_contributor_max = _optional_max( + self.effective_contributor_max, other.effective_contributor_max + ) + if other.weight_sweep is not None: + if self.weight_sweep is None: + self.weight_sweep = CWBWeightSweepDiagnostics() + self.weight_sweep.absorb(other.weight_sweep) + + def record_weighting( + self, + similarities: torch.Tensor, + accepted: torch.Tensor, + weights: torch.Tensor, + ) -> None: + count = similarities.numel() + self.weighting_contributors += count + self.accepted_weighting_contributors += int(accepted.sum().item()) + minimum = float(similarities.min().item()) + maximum = float(similarities.max().item()) + self.consensus_similarity_sum += float(similarities.sum().item()) + self.consensus_similarity_min = _optional_min( + self.consensus_similarity_min, minimum + ) + self.consensus_similarity_max = _optional_max( + self.consensus_similarity_max, maximum + ) + dominant = float(weights.max().item()) + effective = float((1.0 / weights.square().sum().clamp_min(1e-12)).item()) + self.dominant_weight_sum += dominant + self.dominant_weight_min = _optional_min(self.dominant_weight_min, dominant) + self.dominant_weight_max = _optional_max(self.dominant_weight_max, dominant) + self.effective_contributor_sum += effective + self.effective_contributor_min = _optional_min( + self.effective_contributor_min, effective + ) + self.effective_contributor_max = _optional_max( + self.effective_contributor_max, effective + ) + + def render(self, preset_name: str, input_count: int, custom: bool) -> str: + source = "connected custom configuration" if custom else preset_name + lines = [ + "CWB MERGE SUMMARY", + f"Settings: {source}", + f"Input contributors: {input_count}", + f"Alignment candidates: {self.alignment_candidates}", + f"Alignment matches: {self.alignment_matches}", + f"Anchor-only groups: {self.anchor_only_groups}", + f"Weighted vector groups: {self.weighting_groups}", + f"All-rejected equal fallbacks: {self.all_rejected_fallbacks}", + f"Zero-weight equal fallbacks: {self.zero_weight_fallbacks}", + "Accepted weighting contributors: " + f"{self.accepted_weighting_contributors}/{self.weighting_contributors}", + "Consensus similarity min/mean/max: " + f"{_format_metric(self.consensus_similarity_min)} / " + f"{_format_metric(self.consensus_similarity_sum / self.weighting_contributors if self.weighting_contributors else None)} / " + f"{_format_metric(self.consensus_similarity_max)}", + "Dominant normalized weight min/mean/max: " + f"{_format_metric(self.dominant_weight_min)} / " + f"{_format_metric(self.dominant_weight_sum / self.weighting_groups if self.weighting_groups else None)} / " + f"{_format_metric(self.dominant_weight_max)}", + "Effective contributors min/mean/max: " + f"{_format_metric(self.effective_contributor_min)} / " + f"{_format_metric(self.effective_contributor_sum / self.weighting_groups if self.weighting_groups else None)} / " + f"{_format_metric(self.effective_contributor_max)}", + f"CUDA OOM CPU fallbacks: {self.cpu_fallbacks}", + ] + if self.lora_groups: + lines.extend([ + f"LoRA groups: {self.lora_groups}", + f"Mixed-rank LoRA groups: {self.mixed_rank_lora_groups}", + f"Genuine LoRA components: {self.genuine_lora_components}", + "Structural rank slots excluded: " + f"{self.structural_rank_slots_excluded}", + f"Explicit mismatch-zero contributors: {self.explicit_zero_contributors}", + ]) + if self.lora_layer_reports: + lines.extend(_render_lora_dimension_groups(self.lora_layer_reports)) + if self.weight_sweep is not None: + lines.append(self.weight_sweep.render()) + return "\n".join(lines) + + +@dataclass +class LoRALayerReport: + layer: str + ranks: tuple[int, ...] + reference_input: int + reference_rank: int + down_dimensions: tuple[int, ...] + up_dimensions: tuple[int, ...] + alignment_method: str + structural_slots_excluded: int + explicit_zero_contributors: int + reference_components: int = 0 + prefix_components: int = 0 + alignable_source_components: int = 0 + matches: int = 0 + anchor_only_components: int = 0 + candidate_score_count: int = 0 + candidate_score_sum: float = 0.0 + candidate_score_min: float | None = None + candidate_score_max: float | None = None + matched_scores: list[float] = field(default_factory=list) + norm_ratios: list[float] = field(default_factory=list) + matrix_shapes: list[tuple[int, int]] = field(default_factory=list) + + def record_candidates(self, scores: torch.Tensor) -> None: + values = scores.detach().float() + count = values.numel() + if not count: + return + minimum = float(values.min().item()) + maximum = float(values.max().item()) + self.candidate_score_count += count + self.candidate_score_sum += float(values.sum().item()) + self.candidate_score_min = ( + minimum if self.candidate_score_min is None + else min(self.candidate_score_min, minimum) + ) + self.candidate_score_max = ( + maximum if self.candidate_score_max is None + else max(self.candidate_score_max, maximum) + ) + + +def _percentile(values: list[float], fraction: float) -> float | None: + if not values: + return None + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _format_metric(value: float | None) -> str: + return "n/a" if value is None else f"{value:.6g}" + + +def _render_lora_dimension_groups( + layers: list[LoRALayerReport], +) -> list[str]: + grouped = {} + for layer in layers: + key = ( + layer.down_dimensions, + layer.up_dimensions, + layer.ranks, + layer.reference_input, + layer.alignment_method, + ) + grouped.setdefault(key, []).append(layer) + + lines = ["", "CWB LORA DIMENSION GROUPS"] + ordered_groups = sorted( + grouped.values(), + key=lambda group: (-len(group), group[0].layer), + ) + for number, group in enumerate(ordered_groups, start=1): + sample = group[0] + candidate_count = sum(item.candidate_score_count for item in group) + candidate_sum = sum(item.candidate_score_sum for item in group) + candidate_min = min( + (item.candidate_score_min for item in group if item.candidate_score_min is not None), + default=None, + ) + candidate_max = max( + (item.candidate_score_max for item in group if item.candidate_score_max is not None), + default=None, + ) + matched_scores = [score for item in group for score in item.matched_scores] + norm_ratios = [ratio for item in group for ratio in item.norm_ratios] + alignable = sum(item.alignable_source_components for item in group) + matches = sum(item.matches for item in group) + reference_components = sum(item.reference_components for item in group) + anchor_only = sum(item.anchor_only_components for item in group) + matrix_shapes = sorted({shape for item in group for shape in item.matrix_shapes}) + matrix_text = ( + ", ".join(f"{rows}x{columns}" for rows, columns in matrix_shapes) + if matrix_shapes else "index alignment" + ) + lines.extend([ + "", + f"Group {number}: {len(group)} layer(s)", + f"Down non-rank dimensions: {sample.down_dimensions}", + f"Up non-rank dimensions: {sample.up_dimensions}", + f"Input ranks: {sample.ranks}", + f"Reference: input {sample.reference_input}, rank {sample.reference_rank}", + f"Alignment: {sample.alignment_method}; matrices: {matrix_text}", + f"Matched source components: {matches}/{alignable} " + f"({(100.0 * matches / alignable) if alignable else 100.0:.2f}%)", + f"Reference-only components: {anchor_only}/{reference_components}", + "Structural rank slots excluded: " + f"{sum(item.structural_slots_excluded for item in group)}", + "Explicit mismatch-zero contributors: " + f"{sum(item.explicit_zero_contributors for item in group)}", + f"Candidate similarity count: {candidate_count}", + "Candidate similarity min/mean/max: " + f"{_format_metric(candidate_min)} / " + f"{_format_metric(candidate_sum / candidate_count if candidate_count else None)} / " + f"{_format_metric(candidate_max)}", + "Matched similarity p05/median/p95: " + f"{_format_metric(_percentile(matched_scores, 0.05))} / " + f"{_format_metric(_percentile(matched_scores, 0.5))} / " + f"{_format_metric(_percentile(matched_scores, 0.95))}", + "Rank-one delta norm ratio count/min/mean/max: " + f"{len(norm_ratios)} / " + f"{_format_metric(min(norm_ratios) if norm_ratios else None)} / " + f"{_format_metric(sum(norm_ratios) / len(norm_ratios) if norm_ratios else None)} / " + f"{_format_metric(max(norm_ratios) if norm_ratios else None)}", + ]) + + def coverage(item: LoRALayerReport) -> float: + return ( + item.matches / item.alignable_source_components + if item.alignable_source_components else 1.0 + ) + + coverage_exceptions = sorted(layers, key=lambda item: (coverage(item), item.layer))[:5] + norm_exceptions = sorted( + (item for item in layers if item.norm_ratios), + key=lambda item: max(abs(value - 1.0) for value in item.norm_ratios), + reverse=True, + )[:5] + lines.extend(["", "CWB LORA EXCEPTIONS", "Lowest match coverage:"]) + lines.extend( + f"- {item.layer}: {item.matches}/{item.alignable_source_components} " + f"({coverage(item) * 100.0:.2f}%)" + for item in coverage_exceptions + ) + lines.append("Largest rank-one delta norm change:") + lines.extend( + f"- {item.layer}: min={min(item.norm_ratios):.6g}, " + f"max={max(item.norm_ratios):.6g}" + for item in norm_exceptions + ) + return lines + + +@dataclass +class LoRAPairSource: + """One genuine normalized pair or one intentional mismatch-zero input.""" + + source_index: int + down: torch.Tensor | None + up: torch.Tensor | None + scale: float = 1.0 + synthetic_zero: bool = False + + @classmethod + def zero(cls, source_index: int) -> "LoRAPairSource": + return cls(source_index, None, None, synthetic_zero=True) + + @property + def rank(self) -> int: + if self.synthetic_zero or self.down is None: + raise ValueError("A synthetic-zero LoRA source has no genuine rank.") + return int(self.down.shape[0]) + + +CWB_CONFIG = io.Custom("CWB_CONFIG") + + +def resolve_cwb_settings( + preset_name: str, + presets: dict[str, dict], + custom_config: CWBSettings | None = None, +) -> CWBSettings: + """Resolve one complete use-case preset or an explicitly connected config.""" + if custom_config is not None: + if not isinstance(custom_config, CWBSettings): + raise TypeError("CWB config input must come from CWB Custom Configuration.") + return custom_config + try: + resolved = presets[preset_name].copy() + except KeyError as exc: + raise ValueError( + f"Unsupported CWB preset '{preset_name}' for this merge type. " + "Select a current preset or connect CWB Custom Configuration." + ) from exc 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 build_custom_cwb_settings(**values) -> CWBSettings: + settings = CWBSettings( + consensus_type=values["consensus_type"], + alignment_method=values["alignment_method"], + alignment_threshold=float(values["alignment_threshold"]), + similarity_threshold=float(values["similarity_threshold"]), + power_alpha=float(values["power_alpha"]), + diversity_beta=float(values["diversity_beta"]), + rescale_norm=bool(values["rescale_norm"]), + global_scale=float(values["global_scale"]), + dynamic_similarity_contrast=bool(values["dynamic_similarity_contrast"]), + soft_comfort_bandpass=bool(values["soft_comfort_bandpass"]), + position_weight=float(values["position_weight"]), + preserve_common_prefix=bool(values["preserve_common_prefix"]), + ) + if not 0.0 <= settings.position_weight <= 1.0: + raise ValueError("Position weight must be between 0.0 and 1.0.") + return settings + + def _common_prefix_length(tensors: list[torch.Tensor]) -> int: if not tensors: return 0 @@ -214,11 +791,16 @@ def merge_consensus_group( settings: CWBSettings, *, apply_global_scale: bool = True, + diagnostics: CWBDiagnostics | None = None, ) -> torch.Tensor: """Merge one aligned vector group using CWB.""" if stacked.shape[0] == 1: merged = stacked[0].clone() else: + if diagnostics is not None: + diagnostics.weighting_groups += 1 + if diagnostics.weight_sweep is not None: + diagnostics.weight_sweep.record(stacked) consensus = ( torch.median(stacked, dim=0).values if settings.consensus_type == "median" @@ -251,9 +833,15 @@ def merge_consensus_group( if weight_sum > 0: row_weights /= weight_sum else: + if diagnostics is not None: + diagnostics.zero_weight_fallbacks += 1 row_weights.fill_(1.0 / len(similarities)) else: + if diagnostics is not None: + diagnostics.all_rejected_fallbacks += 1 row_weights.fill_(1.0 / len(similarities)) + if diagnostics is not None: + diagnostics.record_weighting(similarities, mask, row_weights) merged = (stacked * row_weights.unsqueeze(1)).sum(dim=0) if settings.rescale_norm: @@ -269,11 +857,14 @@ def merge_consensus_group( def _greedy_similarity_matches( similarities: torch.Tensor, settings: CWBSettings, + diagnostics: CWBDiagnostics | None = None, ) -> list[int]: """Match source rows to reference rows without duplicating the score matrix.""" scores = _position_biased_scores(similarities, settings.position_weight) scores.masked_fill_(similarities < settings.alignment_threshold, -100.0) matched = [-1] * similarities.shape[0] + if diagnostics is not None: + diagnostics.alignment_candidates += similarities.shape[0] for _ in range(min(similarities.shape)): flat_index = torch.argmax(scores) best = scores.flatten()[flat_index].item() @@ -282,6 +873,8 @@ def _greedy_similarity_matches( ref_row = int((flat_index // similarities.shape[1]).item()) source_row = int((flat_index % similarities.shape[1]).item()) matched[ref_row] = source_row + if diagnostics is not None: + diagnostics.alignment_matches += 1 scores[ref_row, :] = -100.0 scores[:, source_row] = -100.0 return matched @@ -293,6 +886,7 @@ def merge_cwb_tensors( *, reference_index: int = 0, allow_similarity_alignment: bool = True, + diagnostics: CWBDiagnostics | None = None, ) -> torch.Tensor: """CWB tensors with a shared rank and trailing vector dimensions.""" if not tensors: @@ -354,15 +948,21 @@ def merge_cwb_tensors( F.normalize(reference_flat, p=2, dim=1, eps=1e-8), F.normalize(flat, p=2, dim=1, eps=1e-8).t(), ) - matched = _greedy_similarity_matches(similarities, settings) + matched = _greedy_similarity_matches(similarities, settings, diagnostics) for ref_row, source_row in enumerate(matched): if source_row >= 0: groups[ref_row].append(tensor[source_row]) merged_rows = [] + if diagnostics is not None and len(tensors) > 1: + diagnostics.anchor_only_groups += sum(len(group) == 1 for group in groups) 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_rows.append( + merge_consensus_group( + stacked, settings, diagnostics=diagnostics + ).reshape(group[0].shape) + ) merged = torch.stack(merged_rows) if prefix_length: merged = torch.cat([preserved_prefix, merged], dim=0) @@ -373,22 +973,36 @@ def _common_lora_prefix_length( downs: list[torch.Tensor], ups: list[torch.Tensor], reference_index: int, + synthetic_zero_count: int = 0, ) -> int: reference_down = downs[reference_index] reference_up = ups[reference_index] - limit = reference_down.shape[0] + limit = min(down.shape[0] for down in downs) common = torch.ones(limit, dtype=torch.bool, device=reference_down.device) for down, up in zip(downs, ups): down_equal = torch.isclose( - reference_down, down, rtol=1e-5, atol=1e-6 + reference_down[:limit], down[:limit], rtol=1e-5, atol=1e-6 ).reshape(limit, -1).all(dim=1) up_equal = torch.isclose( - reference_up.movedim(1, 0), - up.movedim(1, 0), + reference_up.movedim(1, 0)[:limit], + up.movedim(1, 0)[:limit], rtol=1e-5, atol=1e-6, ).reshape(limit, -1).all(dim=1) common &= down_equal & up_equal + if synthetic_zero_count: + common &= torch.isclose( + reference_down[:limit], + torch.zeros((), device=reference_down.device, dtype=reference_down.dtype), + rtol=1e-5, + atol=1e-6, + ).reshape(limit, -1).all(dim=1) + common &= torch.isclose( + reference_up.movedim(1, 0)[:limit], + torch.zeros((), device=reference_up.device, dtype=reference_up.dtype), + rtol=1e-5, + atol=1e-6, + ).reshape(limit, -1).all(dim=1) mismatch = torch.nonzero(~common, as_tuple=False) return limit if mismatch.numel() == 0 else int(mismatch[0].item()) @@ -399,23 +1013,45 @@ def merge_cwb_lora_pairs( settings: CWBSettings, *, reference_index: int, + synthetic_zero_count: int = 0, + diagnostics: CWBDiagnostics | None = None, + layer_report: LoRALayerReport | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """CWB LoRA rank components while keeping A rows paired with B columns.""" if not downs or len(downs) != len(ups): raise ValueError("CWB requires matching LoRA A/B source lists.") - rank = downs[reference_index].shape[0] - if any(down.shape[0] != rank or up.shape[1] != rank for down, up in zip(downs, ups)): - raise ValueError("CWB LoRA sources must be padded to one shared rank.") - - prefix = ( - _common_lora_prefix_length(downs, ups, reference_index) - if settings.preserve_common_prefix else 0 - ) + if not 0 <= reference_index < len(downs): + raise ValueError("CWB LoRA reference index is out of range.") + for down, up in zip(downs, ups): + if down.ndim < 2 or up.ndim < 2 or down.shape[0] <= 0: + raise ValueError("CWB LoRA sources must contain positive-rank factor pairs.") + if down.shape[0] != up.shape[1]: + raise ValueError("CWB LoRA down/up factors must have matching ranks.") reference_down = downs[reference_index] reference_up = ups[reference_index] + rank = reference_down.shape[0] + if any( + tuple(down.shape[1:]) != tuple(reference_down.shape[1:]) + or up.shape[0] != reference_up.shape[0] + or tuple(up.shape[2:]) != tuple(reference_up.shape[2:]) + for down, up in zip(downs, ups) + ): + raise ValueError("CWB LoRA sources must have matching non-rank dimensions.") + + prefix = ( + _common_lora_prefix_length( + downs, ups, reference_index, synthetic_zero_count + ) + if settings.preserve_common_prefix else 0 + ) + if layer_report is not None: + layer_report.prefix_components = prefix + layer_report.reference_components = rank - prefix reference_up_components = reference_up.movedim(1, 0) down_groups = [[reference_down[row]] for row in range(prefix, rank)] up_groups = [[reference_up_components[row]] for row in range(prefix, rank)] + matched_score_tensors = [] + norm_ratio_tensors = [] ref_down_body = reference_down[prefix:].reshape(rank - prefix, -1) ref_up_body = reference_up_components[prefix:].reshape(rank - prefix, -1) @@ -424,8 +1060,12 @@ def merge_cwb_lora_pairs( continue source_down_components = down[prefix:] source_up_components = up.movedim(1, 0)[prefix:] - source_down = source_down_components.reshape(rank - prefix, -1) - source_up = source_up_components.reshape(rank - prefix, -1) + if source_down_components.shape[0] == 0: + continue + if layer_report is not None: + layer_report.alignable_source_components += source_down_components.shape[0] + source_down = source_down_components.reshape(source_down_components.shape[0], -1) + source_up = source_up_components.reshape(source_up_components.shape[0], -1) if settings.alignment_method == "index": matched = list(range(source_down.shape[0])) down_similarities = up_similarities = None @@ -439,11 +1079,23 @@ def merge_cwb_lora_pairs( F.normalize(source_up, p=2, dim=1, eps=1e-8).T, ) contribution_similarities = down_similarities * up_similarities - matched = _greedy_similarity_matches(contribution_similarities, settings) + if layer_report is not None: + layer_report.matrix_shapes.append(tuple(contribution_similarities.shape)) + layer_report.record_candidates(contribution_similarities) + matched = _greedy_similarity_matches( + contribution_similarities, settings, diagnostics + ) for ref_row, source_row in enumerate(matched): if source_row < 0 or ref_row >= len(down_groups): continue + if layer_report is not None: + layer_report.matches += 1 + if down_similarities is not None: + matched_score_tensors.append( + down_similarities[ref_row, source_row] + * up_similarities[ref_row, source_row] + ) source_down_row = source_down_components[source_row] source_up_column = source_up_components[source_row] if ( @@ -456,20 +1108,63 @@ def merge_cwb_lora_pairs( down_groups[ref_row].append(source_down_row) up_groups[ref_row].append(source_up_column) - merged_down_rows = [ - merge_consensus_group( - torch.stack([component.reshape(-1) for component in group]), + if synthetic_zero_count: + zero_down = torch.zeros_like(reference_down[0]) + zero_up = torch.zeros_like(reference_up_components[0]) + for down_group, up_group in zip(down_groups, up_groups): + down_group.extend([zero_down] * synthetic_zero_count) + up_group.extend([zero_up] * synthetic_zero_count) + + if diagnostics is not None and len(downs) > 1: + diagnostics.anchor_only_groups += sum( + len(group) == 1 for group in down_groups + ) + if layer_report is not None: + layer_report.anchor_only_components = sum( + len(group) == 1 for group in down_groups + ) + merged_down_rows = [] + merged_up_columns = [] + for down_group, up_group in zip(down_groups, up_groups): + merged_down = merge_consensus_group( + torch.stack([component.reshape(-1) for component in down_group]), settings, apply_global_scale=False, - ).reshape(group[0].shape) - for group in down_groups - ] - merged_up_columns = [ - merge_consensus_group( - torch.stack([component.reshape(-1) for component in group]), settings - ).reshape(group[0].shape) - for group in up_groups - ] + diagnostics=diagnostics, + ).reshape(down_group[0].shape) + merged_up = merge_consensus_group( + torch.stack([component.reshape(-1) for component in up_group]), + settings, + diagnostics=diagnostics, + ).reshape(up_group[0].shape) + merged_down_rows.append(merged_down) + merged_up_columns.append(merged_up) + if layer_report is not None: + input_norms = torch.stack([ + torch.linalg.vector_norm(down_component) + * torch.linalg.vector_norm(up_component) + for down_component, up_component in zip(down_group, up_group) + ]) + mean_input_norm = input_norms.mean() + output_norm = ( + torch.linalg.vector_norm(merged_down) + * torch.linalg.vector_norm(merged_up) + ) + norm_ratio_tensors.append( + torch.where( + mean_input_norm > 0.0, + output_norm / mean_input_norm.clamp_min(1e-12), + torch.full_like(output_norm, torch.nan), + ) + ) + if layer_report is not None: + if matched_score_tensors: + layer_report.matched_scores.extend( + torch.stack(matched_score_tensors).detach().float().cpu().tolist() + ) + if norm_ratio_tensors: + ratios = torch.stack(norm_ratio_tensors).detach().float().cpu().tolist() + layer_report.norm_ratios.extend(value for value in ratios if value == value) if prefix: merged_down_rows = [*reference_down[:prefix], *merged_down_rows] merged_up_columns = [*reference_up_components[:prefix], *merged_up_columns] @@ -511,6 +1206,7 @@ def _merge_tensors_to_cpu( reference_index: int = 0, allow_similarity_alignment: bool, operation_label: str = "tensor", + diagnostics: CWBDiagnostics | None = None, ) -> torch.Tensor: """Own one tensor operation's GPU lifetime and return only its CPU result.""" def execute(target_device: str) -> torch.Tensor: @@ -521,6 +1217,7 @@ def _merge_tensors_to_cpu( settings, reference_index=reference_index, allow_similarity_alignment=allow_similarity_alignment, + diagnostics=diagnostics, ) return merged.to(target_dtype).cpu().contiguous() finally: @@ -536,45 +1233,91 @@ def _merge_tensors_to_cpu( "[CWB Merge] CUDA OOM for '%s'; retrying this layer on CPU.", operation_label, ) + if diagnostics is not None: + diagnostics.cpu_fallbacks += 1 return execute("cpu") def _merge_lora_pair_to_cpu( - pair_sources: list[tuple[int, torch.Tensor, torch.Tensor, float]], + pair_sources: list[LoRAPairSource], settings: CWBSettings, device: str, target_dtype: torch.dtype, operation_label: str = "LoRA pair", + diagnostics: CWBDiagnostics | None = None, ) -> tuple[torch.Tensor, torch.Tensor, int]: """Own one LoRA pair operation's GPU lifetime and return CPU A/B tensors.""" - max_rank = max(down.shape[0] for _, down, _, _ in pair_sources) + genuine_sources = [source for source in pair_sources if not source.synthetic_zero] + if not genuine_sources: + raise ValueError("CWB LoRA merge requires at least one genuine factor pair.") + max_rank = max(source.rank for source in genuine_sources) reference_index = max( - range(len(pair_sources)), - key=lambda index: pair_sources[index][1].shape[0], + range(len(genuine_sources)), + key=lambda index: genuine_sources[index].rank, ) + reference_source = genuine_sources[reference_index] + if reference_source.down is None or reference_source.up is None: + raise ValueError("CWB LoRA reference source is missing its factors.") + synthetic_zero_count = len(pair_sources) - len(genuine_sources) + ranks = [source.rank for source in genuine_sources] + structural_slots_excluded = sum(max_rank - rank for rank in ranks) + if diagnostics is not None: + diagnostics.lora_groups += 1 + diagnostics.mixed_rank_lora_groups += int(len(set(ranks)) > 1) + diagnostics.genuine_lora_components += sum(ranks) + diagnostics.structural_rank_slots_excluded += structural_slots_excluded + diagnostics.explicit_zero_contributors += synthetic_zero_count + def execute(target_device: str) -> tuple[torch.Tensor, torch.Tensor, int]: - downs = [ - _to_compute(_pad_rank_axis(down, 0, max_rank), target_device) - for _, down, _, _ in pair_sources - ] - ups = [ - _to_compute(_pad_rank_axis(up, 1, max_rank), target_device) * scale - for _, _, up, scale in pair_sources - ] + downs = [] + ups = [] + attempt_diagnostics = ( + CWBDiagnostics( + weight_sweep=( + CWBWeightSweepDiagnostics() + if diagnostics.weight_sweep is not None else None + ) + ) + if diagnostics is not None else None + ) + layer_report = LoRALayerReport( + layer=operation_label, + ranks=tuple(ranks), + reference_input=reference_source.source_index + 1, + reference_rank=max_rank, + down_dimensions=tuple(reference_source.down.shape[1:]), + up_dimensions=( + int(reference_source.up.shape[0]), + *tuple(reference_source.up.shape[2:]), + ), + alignment_method=settings.alignment_method, + structural_slots_excluded=structural_slots_excluded, + explicit_zero_contributors=synthetic_zero_count, + ) try: + for source in genuine_sources: + if source.down is None or source.up is None: + raise ValueError("A genuine LoRA source is missing its factors.") + downs.append(_to_compute(source.down, target_device)) + ups.append(_to_compute(source.up, target_device) * source.scale) merged_down, merged_up = merge_cwb_lora_pairs( downs, ups, settings, reference_index=reference_index, + synthetic_zero_count=synthetic_zero_count, + diagnostics=attempt_diagnostics, + layer_report=layer_report, ) - return ( - merged_down.to(target_dtype).cpu().contiguous(), - merged_up.to(target_dtype).cpu().contiguous(), - max_rank, - ) + output_down = merged_down.to(target_dtype).cpu().contiguous() + output_up = merged_up.to(target_dtype).cpu().contiguous() + if diagnostics is not None and attempt_diagnostics is not None: + diagnostics.absorb_runtime(attempt_diagnostics) + diagnostics.lora_layer_reports.append(layer_report) + return output_down, output_up, max_rank finally: - del downs, ups + downs.clear() + ups.clear() try: return execute(device) @@ -586,6 +1329,8 @@ def _merge_lora_pair_to_cpu( "[CWB Merge] CUDA OOM for '%s'; retrying this layer on CPU.", operation_label, ) + if diagnostics is not None: + diagnostics.cpu_fallbacks += 1 return execute("cpu") @@ -647,15 +1392,6 @@ def _secondary_lora_map( 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.""" @@ -668,7 +1404,14 @@ class ConsensusMergerLogic: *, embedding_union: bool = False, lora_mode: bool = False, - ) -> str: + ) -> str | tuple[str, str]: + diagnostics = CWBDiagnostics( + weight_sweep=( + CWBWeightSweepDiagnostics() + if params.get("counterfactual_weight_sweep", False) else None + ) + ) + params = {**params, "_cwb_diagnostics": diagnostics} paths = [] for name in model_names: path = folder_paths.get_full_path(model_type, name) @@ -691,17 +1434,34 @@ class ConsensusMergerLogic: ) 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, + settings = resolve_cwb_settings( + params["cwb_preset"], + _presets_for( + embedding_union=embedding_union, + lora_mode=lora_mode, + ), + params.get("cwb_config"), ) + if lora_mode: + result = cls._merge_loras( + handlers, low_bit_sets, model_type, params, settings + ) + else: + result = cls._merge_generic( + handlers, + low_bit_sets, + model_type, + params, + settings, + embedding_union=embedding_union, + ) + if params.get("_return_cwb_report"): + return result, diagnostics.render( + params["cwb_preset"], + len(model_names), + params.get("cwb_config") is not None, + ) + return result finally: for handler in handlers: handler.__exit__(None, None, None) @@ -709,9 +1469,12 @@ class ConsensusMergerLogic: @staticmethod def _output_path(model_type: str, output_filename: str) -> str: - output_dir = os.path.join(folder_paths.models_dir, model_type) - os.makedirs(output_dir, exist_ok=True) - return os.path.join(output_dir, f"{output_filename.strip()}.safetensors") + return canonical_model_artifact_path(model_type, output_filename)[0] + + @staticmethod + def _output_name(model_type: str, output_path: str) -> str: + output_dir = os.path.abspath(os.path.join(folder_paths.models_dir, model_type)) + return os.path.relpath(output_path, output_dir).replace(os.sep, "/") @classmethod def _merge_generic( @@ -892,6 +1655,7 @@ class ConsensusMergerLogic: reference_index=reference_index, allow_similarity_alignment=embedding_union, operation_label=key, + diagnostics=params.get("_cwb_diagnostics"), ) writer.write_batch([(key, merged)]) del merged @@ -904,7 +1668,7 @@ class ConsensusMergerLogic: secondary_only_copied, secondary_only_merged, ) - return os.path.basename(output_path) + return cls._output_name(model_type, output_path) @classmethod def _merge_loras(cls, handlers, low_bit_sets, model_type, params, settings): @@ -1109,7 +1873,6 @@ class ConsensusMergerLogic: compatible = [] direct = [] down = up = alpha = tensor = anchor_tensor = None - anchor_down_template = anchor_up_template = None matches = core_matches(core) anchor_index, anchor_block, anchor_keys = next( (index, block, keys) @@ -1175,8 +1938,6 @@ class ConsensusMergerLogic: if "down" in anchor_keys and "up" in anchor_keys: pair_sources = [] pair_failed = False - anchor_down_template = current_loaded[(anchor_index, anchor_keys["down"])] - anchor_up_template = current_loaded[(anchor_index, anchor_keys["up"])] for index, _, keys in matches: if not keys: if secondary_only: @@ -1186,12 +1947,7 @@ class ConsensusMergerLogic: if mismatch_mode == "skip": pair_failed = True break - pair_sources.append(( - index, - torch.zeros_like(anchor_down_template), - torch.zeros_like(anchor_up_template), - 1.0, - )) + pair_sources.append(LoRAPairSource.zero(index)) continue if "down" not in keys or "up" not in keys: if mismatch_mode == "error": @@ -1199,12 +1955,7 @@ class ConsensusMergerLogic: if mismatch_mode == "skip": pair_failed = True break - pair_sources.append(( - index, - torch.zeros_like(anchor_down_template), - torch.zeros_like(anchor_up_template), - 1.0, - )) + pair_sources.append(LoRAPairSource.zero(index)) continue down = current_loaded[(index, keys["down"])] up = current_loaded[(index, keys["up"])] @@ -1214,12 +1965,7 @@ class ConsensusMergerLogic: if mismatch_mode == "skip": pair_failed = True break - pair_sources.append(( - index, - torch.zeros_like(anchor_down_template), - torch.zeros_like(anchor_up_template), - 1.0, - )) + pair_sources.append(LoRAPairSource.zero(index)) continue alpha = ( current_loaded[(index, keys["alpha"])] @@ -1230,7 +1976,9 @@ class ConsensusMergerLogic: ) current_loaded[(index, keys["down"])] = down current_loaded[(index, keys["up"])] = up - pair_sources.append((index, down, up, 1.0)) + pair_sources.append( + LoRAPairSource(index, down, up) + ) if pair_failed: preserve_roles( writer, @@ -1240,10 +1988,28 @@ class ConsensusMergerLogic: anchor_index, ) elif pair_sources: - anchor_down = pair_sources[0][1] - anchor_up = pair_sources[0][2] + anchor_source = next( + source for source in pair_sources + if not source.synthetic_zero + ) + anchor_down = anchor_source.down + anchor_up = anchor_source.up + if anchor_down is None or anchor_up is None: + raise ValueError( + f"LoRA pair '{anchor_block}' has no genuine anchor." + ) compatible = [] - for index, down, up, scale in pair_sources: + for source in pair_sources: + if source.synthetic_zero: + compatible.append(source) + continue + index = source.source_index + down = source.down + up = source.up + if down is None or up is None: + raise ValueError( + f"LoRA pair '{anchor_block}' has missing factors." + ) valid = ( tuple(down.shape[1:]) == tuple(anchor_down.shape[1:]) and up.shape[0] == anchor_up.shape[0] @@ -1255,14 +2021,9 @@ class ConsensusMergerLogic: if mismatch_mode == "skip": compatible = [] break - compatible.append(( - index, - torch.zeros_like(anchor_down), - torch.zeros_like(anchor_up), - 1.0, - )) + compatible.append(LoRAPairSource.zero(index)) continue - compatible.append((index, down, up, scale)) + compatible.append(source) if not compatible: preserve_roles( writer, @@ -1283,6 +2044,7 @@ class ConsensusMergerLogic: params["process_device"], target_dtype, operation_label=anchor_block, + diagnostics=params.get("_cwb_diagnostics"), ) writer.write( canonical_lora_key(anchor_block, "down"), @@ -1385,6 +2147,7 @@ class ConsensusMergerLogic: target_dtype, allow_similarity_alignment=False, operation_label=output_key, + diagnostics=params.get("_cwb_diagnostics"), ) writer.write(output_key, merged) del merged @@ -1402,7 +2165,6 @@ class ConsensusMergerLogic: direct.clear() current_loaded.clear() down = up = alpha = tensor = anchor_tensor = None - anchor_down_template = anchor_up_template = None pbar.update(1) if preserved_companion_groups: @@ -1416,10 +2178,16 @@ class ConsensusMergerLogic: secondary_only_copied, secondary_only_merged, ) - return os.path.basename(output_path) + return cls._output_name(model_type, output_path) -def _common_inputs(model_type: str, input_count: int, default_filename: str): +def _common_inputs( + model_type: str, + input_count: int, + default_filename: str, + presets: dict[str, dict], + default_preset: str, +): label = model_type inputs = [ io.Combo.Input( @@ -1447,89 +2215,14 @@ def _common_inputs(model_type: str, input_count: int, default_filename: str): inputs.extend([ io.Combo.Input( "cwb_preset", - options=CWB_PRESET_OPTIONS, - default="baseline", - tooltip="Select custom to use every manual CWB control. Named presets set their listed controls; position/prefix controls, most alignment thresholds, and any explicit non-1.0 global scale remain active.", + options=list(presets), + default=default_preset, + tooltip="Use-case preset. Name suffixes expose alignment, consensus, norm rescaling, DSC, soft comfort bandpass, and prefix preservation. A connected CWB Config overrides it completely.", ), - io.Combo.Input( - "consensus_type", - options=["mean", "median"], - default="median", - tooltip="Baseline vector used to score contributors. Median resists coordinate outliers; mean represents their arithmetic center. Manual value requires the custom preset.", - ), - io.Combo.Input( - "alignment_method", - options=["index", "similarity"], - default="similarity", - tooltip="Pair rows by original index or greedy cosine similarity. Row reordering applies only to embeddings and LoRA rank components; ordinary model tensors keep fixed coordinates.", - ), - io.Float.Input( - "alignment_threshold", - default=0.4, - min=0.0, - max=1.0, - step=0.01, - tooltip="Minimum original cosine score for a greedy row match. Applies only to similarity-aligned embeddings and LoRA rank components; unmatched rows do not contribute.", - ), - io.Float.Input( - "similarity_threshold", - default=0.0, - min=-1.0, - max=1.0, - step=0.01, - tooltip="Minimum cosine similarity to the mean/median consensus for a paired vector to receive weight. If every vector is rejected, CWB falls back to equal weights.", - ), - io.Float.Input( - "power_alpha", - default=2.0, - min=0.0, - max=10.0, - step=0.1, - tooltip="Exponent applied to accepted non-negative consensus similarities. Higher values concentrate weight on vectors closest to consensus. Manual value requires the custom preset.", - ), - io.Float.Input( - "diversity_beta", - default=0.0, - min=0.0, - max=10.0, - step=0.1, - tooltip="Bandpass exponent that suppresses near-consensus dominance to retain more variation. Zero disables diversity weighting. Manual value requires the custom preset.", - ), - io.Boolean.Input( - "rescale_norm", - default=False, - tooltip="Preserve vector energy after blending: keep the merged direction but set its L2 norm to the mean norm of participating vectors. Applied per row; for LoRA, separately per paired A row and B column.", - ), - io.Float.Input( - "global_scale", - default=1.0, - min=0.0, - max=10.0, - step=0.01, - tooltip="Multiply the final merged contribution. For LoRA it is applied once through B/up, not to both factors. Named presets remain neutral unless this is explicitly set away from 1.0.", - ), - io.Boolean.Input( - "dynamic_similarity_contrast", - default=False, - tooltip="Remap unequal consensus similarities into 0.7-1.0 before alpha/beta weighting to increase relative separation. It does not change row-alignment scores.", - ), - io.Boolean.Input( - "soft_comfort_bandpass", - default=False, - tooltip="When diversity_beta is positive, broaden its distance term from 1.001-similarity to 1.5-similarity. Has no effect when diversity_beta is zero.", - ), - io.Float.Input( - "position_weight", - default=0.0, - min=0.0, - max=1.0, - step=0.01, - tooltip="Blend positional proximity into greedy row selection: 0 uses cosine only; 1 uses position affinity only. Relevant only to similarity-aligned embeddings and LoRA rank components.", - ), - io.Boolean.Input( - "preserve_common_prefix", - default=False, - tooltip="Copy the longest numerically identical leading first-axis span from Model A before CWB. The current merger applies this to all tensors with at least two dimensions.", + CWB_CONFIG.Input( + "cwb_config", + optional=True, + tooltip="Optional settings from CWB Custom Configuration. When connected, it completely overrides the selected preset.", ), io.Combo.Input( "mismatch_mode", @@ -1588,6 +2281,36 @@ def _common_inputs(model_type: str, input_count: int, default_filename: str): return inputs +class CWBCustomConfiguration(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="CWBCustomConfiguration", + display_name="CWB Custom Configuration", + category="ModelUtils/Merging/Configuration", + description="Build complete custom Consensus-Weighted Blending settings for any CWB merge node.", + inputs=[ + io.Combo.Input("consensus_type", options=["mean", "median"], default="median", tooltip="Mean is a symmetric center; median is more robust with three or more contributors."), + io.Combo.Input("alignment_method", options=["index", "similarity"], default="similarity", tooltip="Pair vectors by position or greedy cosine similarity where the merge type supports alignment."), + io.Float.Input("alignment_threshold", default=0.0, min=0.0, max=1.0, step=0.01, tooltip="Minimum cosine product or score for accepting a similarity-aligned row match."), + io.Float.Input("similarity_threshold", default=0.0, min=-1.0, max=1.0, step=0.01, tooltip="Minimum contributor similarity to the group consensus before weighting."), + io.Float.Input("power_alpha", default=2.0, min=0.0, max=10.0, step=0.1, tooltip="Exponent applied to accepted non-negative consensus similarities."), + io.Float.Input("diversity_beta", default=10.0, min=0.0, max=10.0, step=0.1, tooltip="Bandpass exponent used to suppress near-consensus dominance; zero disables it."), + io.Boolean.Input("rescale_norm", default=True, tooltip="Set each merged vector norm to the mean participating norm after weighting."), + io.Float.Input("global_scale", default=1.0, min=0.0, max=10.0, step=0.01, tooltip="Multiply the final contribution once; LoRA applies this through the up factor."), + io.Boolean.Input("dynamic_similarity_contrast", default=False, tooltip="Remap unequal consensus similarities into 0.7 to 1.0 before alpha and beta weighting."), + io.Boolean.Input("soft_comfort_bandpass", default=True, tooltip="Use 1.5 minus similarity instead of 1.001 minus similarity for diversity weighting."), + io.Float.Input("position_weight", default=0.05, min=0.0, max=1.0, step=0.01, tooltip="Blend positional affinity into similarity-based greedy matching."), + io.Boolean.Input("preserve_common_prefix", default=False, tooltip="Copy a numerically identical leading component span from the anchor."), + ], + outputs=[CWB_CONFIG.Output(display_name="cwb_config")], + ) + + @classmethod + def execute(cls, **kwargs): + return io.NodeOutput(build_custom_cwb_settings(**kwargs)) + + class _CWBMergerNode(io.ComfyNode): MODEL_TYPE = "diffusion_models" INPUT_COUNT = 2 @@ -1596,10 +2319,21 @@ class _CWBMergerNode(io.ComfyNode): DEFAULT_FILENAME = "cwb_merged" EMBEDDING_UNION = False LORA_MODE = False + DEFAULT_PRESET = "balanced_mean" @classmethod def define_schema(cls): - inputs = _common_inputs(cls.MODEL_TYPE, cls.INPUT_COUNT, cls.DEFAULT_FILENAME) + presets = _presets_for( + embedding_union=cls.EMBEDDING_UNION, + lora_mode=cls.LORA_MODE, + ) + inputs = _common_inputs( + cls.MODEL_TYPE, + cls.INPUT_COUNT, + cls.DEFAULT_FILENAME, + presets, + cls.DEFAULT_PRESET, + ) if cls.LORA_MODE: inputs.append(io.Boolean.Input( "include_1d_diffs", @@ -1613,8 +2347,9 @@ class _CWBMergerNode(io.ComfyNode): description="Merge two or three safetensors files with Consensus-Weighted Blending. Use each input tooltip or Documentation Only for exact control scope.", inputs=inputs, outputs=[ - io.String.Output(display_name="output_filename"), + io.AnyType.Output(display_name="output_filename"), io.String.Output(display_name="documentation"), + io.String.Output(display_name="cwb_report"), ], is_experimental=True, ) @@ -1623,18 +2358,22 @@ class _CWBMergerNode(io.ComfyNode): 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) + return io.NodeOutput( + "Documentation mode active. No merge performed.", + documentation, + "CWB report unavailable because no merge was performed.", + ) names = [kwargs["model_a"], kwargs["model_b"]] if cls.INPUT_COUNT == 3: names.append(kwargs["model_c"]) - filename = ConsensusMergerLogic.execute( + filename, report = ConsensusMergerLogic.execute( names, cls.MODEL_TYPE, - kwargs, + {**kwargs, "_return_cwb_report": True}, embedding_union=cls.EMBEDDING_UNION, lora_mode=cls.LORA_MODE, ) - return io.NodeOutput(filename, documentation) + return io.NodeOutput(filename, documentation, report) class CWBCheckpointTwoMerger(_CWBMergerNode): @@ -1684,6 +2423,7 @@ class CWBLoRATwoMerger(_CWBMergerNode): DISPLAY_NAME = "CWB Merge LoRAs (2 Models)" DEFAULT_FILENAME = "cwb_merged_2_lora" LORA_MODE = True + DEFAULT_PRESET = "broad_sim_medn_rn_softcb" class CWBLoRAThreeMerger(CWBLoRATwoMerger): @@ -1693,12 +2433,81 @@ class CWBLoRAThreeMerger(CWBLoRATwoMerger): DEFAULT_FILENAME = "cwb_merged_3_lora" +class CWBLoRAMultiMerger(io.ComfyNode): + @classmethod + def define_schema(cls): + lora_options = folder_paths.get_filename_list("loras") + optional_loras = ["None", *lora_options] + return io.Schema( + node_id="CWBLoRAMultiMerger", + display_name="CWB LoRA Multi-Merge", + category="ModelUtils/LoRA/Merge", + description="Consensus-merge 2 to 8 equal-prior LoRAs with bounded UEL streaming.", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"], default="MERGE", tooltip="MERGE writes an output; DOCUMENTATION ONLY opens no model files."), + io.Combo.Input("lora_count", options=[str(i) for i in range(2, 9)], default="2", tooltip="Number of consecutive LoRA inputs to include."), + io.Combo.Input("lora_1", options=lora_options, tooltip="Anchor LoRA and metadata source."), + io.Combo.Input("lora_2", options=lora_options, tooltip="Second equal-prior LoRA contributor."), + io.Combo.Input("lora_3", options=optional_loras, default="None", tooltip="Optional third equal-prior contributor."), + io.Combo.Input("lora_4", options=optional_loras, default="None", tooltip="Optional fourth equal-prior contributor."), + io.Combo.Input("lora_5", options=optional_loras, default="None", tooltip="Optional fifth equal-prior contributor."), + io.Combo.Input("lora_6", options=optional_loras, default="None", tooltip="Optional sixth equal-prior contributor."), + io.Combo.Input("lora_7", options=optional_loras, default="None", tooltip="Optional seventh equal-prior contributor."), + io.Combo.Input("lora_8", options=optional_loras, default="None", tooltip="Optional eighth equal-prior contributor."), + io.Combo.Input("cwb_preset", options=list(LORA_CWB_PRESETS), default="broad_sim_medn_rn_softcb", tooltip="LoRA-specific preset. A connected CWB Config overrides it completely."), + CWB_CONFIG.Input("cwb_config", optional=True, tooltip="Optional complete override from CWB Custom Configuration."), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip", tooltip="Preserve the anchor, insert zeros, or abort for missing and incompatible logical groups."), + io.String.Input("output_filename", default="cwb_merged_multi_lora", tooltip="Filename without extension under ComfyUI's LoRA directory."), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"], default="bf16", tooltip="Requested dtype for generated floating factors."), + io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda", tooltip="Per-layer FP32 processing device; CUDA OOM retries the affected layer on CPU."), + io.String.Input("exclude_patterns", default="", multiline=True, tooltip="Preserve matching logical groups from the anchor."), + io.String.Input("discard_patterns", default="", multiline=True, tooltip="Omit matching tensors or logical groups from the output."), + io.Boolean.Input("glob_patterns", default=False, tooltip="Interpret filter entries as shell-style globs instead of regular expressions."), + io.Boolean.Input("lazy_load", default=True, tooltip="Use bounded UEL work-unit streaming and release each completed input layer."), + io.Boolean.Input("force_clear_cache", default=True, tooltip="Collect Python and CUDA caches before each layer at a potential speed cost."), + io.Boolean.Input("override_dtype", default=False, tooltip="Force generated floating factors to the requested save dtype."), + io.Boolean.Input("include_1d_diffs", default=False, tooltip="CWB-merge 1D direct differences as FP32 instead of preserving the anchor."), + io.Boolean.Input("counterfactual_weight_sweep", default=False, tooltip="Evaluate alpha, beta, similarity-threshold, DSC, and comfort-bandpass weight combinations from each already-computed consensus similarity vector and append them to the report without additional model loads or saved outputs."), + ], + 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.", + ) + count = int(kwargs["lora_count"]) + names = [kwargs[f"lora_{index}"] for index in range(1, count + 1)] + missing = [index for index, name in enumerate(names, start=1) if not name or name == "None"] + if missing: + joined = ", ".join(str(index) for index in missing) + raise ValueError(f"LoRA Count includes unselected input(s): {joined}.") + filename, report = ConsensusMergerLogic.execute( + names, + "loras", + {**kwargs, "_return_cwb_report": True}, + lora_mode=True, + ) + return io.NodeOutput(filename, documentation, report) + + 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 + DEFAULT_PRESET = "balanced_idx_mean" class CWBEmbeddingThreeMerger(CWBEmbeddingTwoMerger): @@ -1709,6 +2518,7 @@ class CWBEmbeddingThreeMerger(CWBEmbeddingTwoMerger): CWB_MERGER_NODES = [ + CWBCustomConfiguration, CWBCheckpointTwoMerger, CWBCheckpointThreeMerger, CWBModelTwoMerger, @@ -1717,6 +2527,7 @@ CWB_MERGER_NODES = [ CWBTextEncoderThreeMerger, CWBLoRATwoMerger, CWBLoRAThreeMerger, + CWBLoRAMultiMerger, CWBEmbeddingTwoMerger, CWBEmbeddingThreeMerger, ] diff --git a/run_cwb_api_workflow.py b/run_cwb_api_workflow.py new file mode 100644 index 0000000..cbaad26 --- /dev/null +++ b/run_cwb_api_workflow.py @@ -0,0 +1,259 @@ +"""Queue the CWB six-LoRA API workflow and print its merge report.""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +import sys +import time +from pathlib import Path +from urllib import error, request + + +DEFAULT_WORKFLOW = Path(__file__).with_name("CWB6LORATEST_API.json") +MERGER_CLASS = "CWBLoRAMultiMerger" +CUSTOM_CLASS = "CWBCustomConfiguration" +PREVIEW_CLASS = "PreviewAny" +CALIBRATION_EXCLUDE_PATTERN = ( + r"^diffusion_model\.blocks\.(?:[12]|[4-9]|1[0-3]|1[5-9]|2[0-3]|2[56])\." + "\n" + r"^diffusion_model\.txtfusion\.layerwise_blocks\.1\." + "\n" + r"^diffusion_model\.txtfusion\.refiner_blocks\.0\." +) + + +def parse_counterfactual_weight_sweep(report: str) -> list[dict]: + pattern = re.compile( + r"^(?Pmean|median) \| " + r"(?P[-+0-9.e]+) \| " + r"(?P[-+0-9.e]+) \| (?P[-+0-9.e]+) \| " + r"(?Ptrue|false) \| (?Ptrue|false) \| " + r"(?P\d+)/(?P\d+) \| (?P\d+) \| " + r"(?P[-+0-9.e]+)/(?P[-+0-9.e]+) \| " + r"(?P[-+0-9.e]+)/(?P[-+0-9.e]+)$", + re.MULTILINE, + ) + rows = [] + for match in pattern.finditer(report): + values = match.groupdict() + rows.append({ + "consensus_type": values["consensus_type"], + "similarity_threshold": float(values["similarity_threshold"]), + "power_alpha": float(values["power_alpha"]), + "diversity_beta": float(values["diversity_beta"]), + "dynamic_similarity_contrast": values["dsc"] == "true", + "soft_comfort_bandpass": values["softcb"] == "true", + "accepted_contributors": int(values["accepted"]), + "contributors": int(values["contributors"]), + "fallbacks": int(values["fallbacks"]), + "dominant_weight_mean": float(values["dominant_mean"]), + "dominant_weight_max": float(values["dominant_max"]), + "effective_contributors_mean": float(values["effective_mean"]), + "effective_contributors_min": float(values["effective_min"]), + }) + return rows + + +def write_counterfactual_csv(path: Path, rows: list[dict]) -> None: + if not rows: + raise RuntimeError("CWB report does not contain a counterfactual weight sweep") + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +def _json_request(url: str, payload=None): + data = None if payload is None else json.dumps(payload).encode("utf-8") + req = request.Request(url, data=data) + if data is not None: + req.add_header("Content-Type", "application/json") + try: + with request.urlopen(req) as response: + return json.load(response) + except error.HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace") + raise RuntimeError(f"ComfyUI request failed ({exc.code}): {detail}") from exc + except error.URLError as exc: + raise RuntimeError(f"Could not reach ComfyUI at {url}: {exc.reason}") from exc + + +def _single_node(workflow: dict, class_type: str) -> tuple[str, dict]: + matches = [ + (node_id, node) + for node_id, node in workflow.items() + if node.get("class_type") == class_type + ] + if len(matches) != 1: + raise RuntimeError( + f"Expected exactly one {class_type} node, found {len(matches)}." + ) + return matches[0] + + +def _linked_preview_nodes(workflow: dict, merger_id: str) -> dict[int, str]: + linked = {} + for node_id, node in workflow.items(): + if node.get("class_type") != PREVIEW_CLASS: + continue + source = node.get("inputs", {}).get("source") + if ( + isinstance(source, list) + and len(source) == 2 + and str(source[0]) == merger_id + and isinstance(source[1], int) + ): + linked[source[1]] = node_id + missing = {0, 2} - linked.keys() + if missing: + raise RuntimeError( + "Workflow must connect merger output 0 (filename) and output 2 " + "(CWB report) to PreviewAny nodes." + ) + return linked + + +def _prepare_workflow( + workflow: dict, + preset: str | None, + output_filename: str | None, + full_blocks: bool = False, +) -> tuple[str, dict[int, str]]: + merger_id, merger = _single_node(workflow, MERGER_CLASS) + _single_node(workflow, CUSTOM_CLASS) + inputs = merger.setdefault("inputs", {}) + + if preset is not None: + inputs["cwb_preset"] = preset + inputs.pop("cwb_config", None) + if output_filename is not None: + inputs["output_filename"] = output_filename + if not full_blocks: + existing = inputs.get("exclude_patterns", "").strip() + inputs["exclude_patterns"] = "\n".join( + value for value in (existing, CALIBRATION_EXCLUDE_PATTERN) if value + ) + + return merger_id, _linked_preview_nodes(workflow, merger_id) + + +def _prompt_payload(workflow: dict) -> dict: + return {"prompt": workflow} + + +def _wait_for_history(server: str, prompt_id: str, poll_seconds: float): + while True: + history = _json_request(f"{server}/history/{prompt_id}") + if prompt_id in history: + return history[prompt_id] + time.sleep(poll_seconds) + + +def _text_output(history_outputs: dict, node_id: str) -> str: + output = history_outputs.get(node_id, {}) + candidates = ( + output.get("text"), + output.get("ui", {}).get("text") if isinstance(output.get("ui"), dict) else None, + ) + for candidate in candidates: + if isinstance(candidate, list) and candidate: + return str(candidate[0]) + if isinstance(candidate, str): + return candidate + raise RuntimeError(f"PreviewAny node {node_id} did not return text in API history.") + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Run the CWB six-LoRA API workflow and collect its report." + ) + parser.add_argument("--workflow", type=Path, default=DEFAULT_WORKFLOW) + parser.add_argument("--server", default="http://127.0.0.1:8188") + parser.add_argument("--poll-seconds", type=float, default=2.0) + parser.add_argument( + "--preset", + help="Disconnect the custom configuration and run this CWB preset.", + ) + parser.add_argument( + "--output-filename", + help="Override the merger output filename without its extension.", + ) + parser.add_argument( + "--result-json", + type=Path, + help="Optionally write prompt ID, filename, and report as JSON.", + ) + parser.add_argument( + "--full-blocks", + action="store_true", + help="Process all blocks for final verification instead of the calibration subset.", + ) + parser.add_argument("--summary-only", action="store_true") + parser.add_argument("--counterfactual-csv", type=Path) + parser.add_argument("--alignment-threshold", type=float) + parser.add_argument("--alignment-method", choices=["index", "similarity"]) + args = parser.parse_args() + + if args.poll_seconds <= 0: + parser.error("--poll-seconds must be greater than zero") + if not args.workflow.is_file(): + parser.error(f"workflow not found: {args.workflow}") + + workflow = json.loads(args.workflow.read_text(encoding="utf-8")) + if args.alignment_threshold is not None: + _, custom = _single_node(workflow, CUSTOM_CLASS) + custom.setdefault("inputs", {})["alignment_threshold"] = ( + args.alignment_threshold + ) + if args.alignment_method is not None: + _, custom = _single_node(workflow, CUSTOM_CLASS) + custom.setdefault("inputs", {})["alignment_method"] = args.alignment_method + _, linked_outputs = _prepare_workflow( + workflow, + args.preset, + args.output_filename, + args.full_blocks, + ) + server = args.server.rstrip("/") + queued = _json_request(f"{server}/prompt", _prompt_payload(workflow)) + prompt_id = queued["prompt_id"] + sys.stdout.write(f"Queued prompt {prompt_id}\n") + + history = _wait_for_history(server, prompt_id, args.poll_seconds) + status = history.get("status", {}) + if status.get("status_str") == "error": + sys.stderr.write(json.dumps(status, indent=2) + "\n") + return 1 + + outputs = history.get("outputs", {}) + output_filename = _text_output(outputs, linked_outputs[0]) + cwb_report = _text_output(outputs, linked_outputs[2]) + result = { + "prompt_id": prompt_id, + "output_filename": output_filename, + "cwb_report": cwb_report, + } + if args.result_json is not None: + args.result_json.parent.mkdir(parents=True, exist_ok=True) + args.result_json.write_text( + json.dumps(result, indent=2) + "\n", + encoding="utf-8", + ) + if args.counterfactual_csv is not None: + write_counterfactual_csv( + args.counterfactual_csv, + parse_counterfactual_weight_sweep(cwb_report), + ) + + sys.stdout.write(f"Output: {output_filename}\n") + if not args.summary_only: + sys.stdout.write(f"\n{cwb_report}\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/run_cwb_preset_sweep.py b/run_cwb_preset_sweep.py new file mode 100644 index 0000000..344a227 --- /dev/null +++ b/run_cwb_preset_sweep.py @@ -0,0 +1,418 @@ +"""Run a resumable one-factor CWB LoRA calibration sweep.""" + +from __future__ import annotations + +import argparse +import csv +import json +import os +import re +import sys +import time +from copy import deepcopy +from pathlib import Path + +import run_cwb_api_workflow as api +import run_model_analysis_api_workflow as analysis_api + + +CONFIG_KEYS = ( + "consensus_type", + "alignment_method", + "alignment_threshold", + "similarity_threshold", + "power_alpha", + "diversity_beta", + "rescale_norm", + "global_scale", + "dynamic_similarity_contrast", + "soft_comfort_bandpass", + "position_weight", + "preserve_common_prefix", +) + +INTEGER_METRICS = { + "alignment_candidates": "Alignment candidates", + "alignment_matches": "Alignment matches", + "anchor_only_groups": "Anchor-only groups", + "weighted_vector_groups": "Weighted vector groups", + "all_rejected_fallbacks": "All-rejected equal fallbacks", + "zero_weight_fallbacks": "Zero-weight equal fallbacks", + "cuda_oom_fallbacks": "CUDA OOM CPU fallbacks", + "lora_groups": "LoRA groups", +} + +TRIPLE_METRICS = { + "consensus_similarity": "Consensus similarity min/mean/max", + "dominant_weight": "Dominant normalized weight min/mean/max", + "effective_contributors": "Effective contributors min/mean/max", +} + +DELTA_METRICS = ( + "alignment_match_fraction", + "anchor_only_groups", + "accepted_weight_fraction", + "consensus_similarity_mean", + "dominant_weight_mean", + "dominant_weight_max", + "effective_contributors_mean", + "effective_contributors_min", + "norm_ratio_mean", + "norm_ratio_min", + "norm_ratio_max", + "all_rejected_fallbacks", + "zero_weight_fallbacks", + "wall_seconds", + "analysis_relative_l2", + "analysis_cosine", + "analysis_exact", + "analysis_norm_ratio", + "analysis_pair_cosine", + "analysis_mean_affinity_a", + "analysis_mean_affinity_b", + "analysis_median_affinity_a", + "analysis_median_affinity_b", +) + + +def _single_custom_node(workflow: dict) -> tuple[str, dict]: + return api._single_node(workflow, api.CUSTOM_CLASS) + + +def _baseline_config(workflow: dict) -> dict: + _, node = _single_custom_node(workflow) + inputs = node.get("inputs", {}) + missing = [key for key in CONFIG_KEYS if key not in inputs] + if missing: + raise RuntimeError( + "Custom configuration is missing required inputs: " + ", ".join(missing) + ) + return {key: inputs[key] for key in CONFIG_KEYS} + + +def build_sweep_cases(baseline: dict) -> list[dict]: + variations = [ + ("alignment_method", "index"), + *[("alignment_threshold", value) for value in (0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01)], + *[("position_weight", value) for value in (0.0, 0.0001, 0.0005, 0.001, 0.005, 0.01)], + ("consensus_type", "mean"), + *[("similarity_threshold", value) for value in (-0.25, 0.1, 0.25, 0.5, 0.75)], + *[("power_alpha", value) for value in (0.5, 1.0, 3.0, 4.0)], + *[("diversity_beta", value) for value in (0.0, 1.0, 2.0, 4.0, 7.0)], + ("rescale_norm", False), + ("dynamic_similarity_contrast", True), + ("soft_comfort_bandpass", False), + ] + cases = [{"id": "baseline", "factor": "baseline", "value": None, "config": baseline.copy()}] + for factor, value in variations: + if baseline[factor] == value: + continue + config = baseline.copy() + config[factor] = value + value_text = str(value).lower().replace(".", "p").replace("-", "neg") + cases.append({ + "id": f"{factor}_{value_text}", + "factor": factor, + "value": value, + "config": config, + }) + return cases + + +def parse_cwb_report(report: str) -> dict: + metrics = {} + for metric, label in INTEGER_METRICS.items(): + match = re.search(rf"^{re.escape(label)}: (\d+)$", report, re.MULTILINE) + if match: + metrics[metric] = int(match.group(1)) + + accepted = re.search( + r"^Accepted weighting contributors: (\d+)/(\d+)$", + report, + re.MULTILINE, + ) + if not accepted: + raise RuntimeError( + "CWB report lacks contributor-weight diagnostics. Restart ComfyUI " + "after loading the current consensus_merger.py before running the sweep." + ) + metrics["accepted_weighting_contributors"] = int(accepted.group(1)) + metrics["weighting_contributors"] = int(accepted.group(2)) + + number = r"(?:n/a|[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:e[-+]?\d+)?)" + for prefix, label in TRIPLE_METRICS.items(): + match = re.search( + rf"^{re.escape(label)}: ({number}) / ({number}) / ({number})$", + report, + re.MULTILINE | re.IGNORECASE, + ) + if not match: + raise RuntimeError(f"CWB report is missing metric: {label}") + for suffix, value in zip(("min", "mean", "max"), match.groups()): + metrics[f"{prefix}_{suffix}"] = None if value == "n/a" else float(value) + + norm_groups = re.findall( + rf"^Rank-one delta norm ratio count/min/mean/max: (\d+) / " + rf"({number}) / ({number}) / ({number})$", + report, + re.MULTILINE | re.IGNORECASE, + ) + numeric_norms = [ + (int(count), *(float(value) for value in (minimum, mean, maximum))) + for count, minimum, mean, maximum in norm_groups + if count != "0" and "n/a" not in (minimum, mean, maximum) + ] + if numeric_norms: + norm_count = sum(group[0] for group in numeric_norms) + metrics["norm_ratio_count"] = norm_count + metrics["norm_ratio_min"] = min(group[1] for group in numeric_norms) + metrics["norm_ratio_mean"] = ( + sum(group[0] * group[2] for group in numeric_norms) / norm_count + ) + metrics["norm_ratio_max"] = max(group[3] for group in numeric_norms) + + candidates = metrics.get("alignment_candidates", 0) + contributors = metrics["weighting_contributors"] + metrics["alignment_match_fraction"] = ( + metrics.get("alignment_matches", 0) / candidates if candidates else 0.0 + ) + metrics["accepted_weight_fraction"] = ( + metrics["accepted_weighting_contributors"] / contributors + if contributors else 0.0 + ) + return metrics + + +def parse_analysis_reports(comparison: str, cwb_report: str) -> dict: + global_section = comparison.split("PARAMETER-WEIGHTED GLOBAL METRICS", 1)[1].split( + "EQUAL-LAYER AVERAGES", 1 + )[0] + + def single(label: str, text: str) -> float: + match = re.search(rf"^{re.escape(label)}: ([-+0-9.e]+)$", text, re.MULTILINE) + if not match: + raise RuntimeError(f"Analysis report is missing metric: {label}") + return float(match.group(1)) + + metrics = { + "analysis_relative_l2": single("relative_l2", global_section), + "analysis_cosine": single("cosine", global_section), + "analysis_exact": single("exact", global_section), + "analysis_norm_ratio": single("norm_ratio", global_section), + } + pairwise = re.search( + r"^Pairwise row cosine mean/min/max: ([-+0-9.e]+) /", + cwb_report, + re.MULTILINE, + ) + mean_affinity = re.search( + r"^Mean-consensus affinity A/B: ([-+0-9.e]+) / ([-+0-9.e]+)$", + cwb_report, + re.MULTILINE, + ) + median_affinity = re.search( + r"^Median-consensus affinity A/B: ([-+0-9.e]+) / ([-+0-9.e]+)$", + cwb_report, + re.MULTILINE, + ) + if not pairwise or not mean_affinity or not median_affinity: + raise RuntimeError("CWB analysis report is missing global affinity metrics") + metrics.update({ + "analysis_pair_cosine": float(pairwise.group(1)), + "analysis_mean_affinity_a": float(mean_affinity.group(1)), + "analysis_mean_affinity_b": float(mean_affinity.group(2)), + "analysis_median_affinity_a": float(median_affinity.group(1)), + "analysis_median_affinity_b": float(median_affinity.group(2)), + }) + return metrics + + +def _run_case( + workflow_template: dict, + case: dict, + server: str, + poll_seconds: float, + full_blocks: bool, + analysis_workflow: dict | None = None, +) -> dict: + workflow = deepcopy(workflow_template) + _, custom = _single_custom_node(workflow) + custom["inputs"].update(case["config"]) + _, linked = api._prepare_workflow( + workflow, + preset=None, + output_filename=None, + full_blocks=full_blocks, + ) + started = time.perf_counter() + queued = api._json_request(f"{server}/prompt", api._prompt_payload(workflow)) + prompt_id = queued["prompt_id"] + history = api._wait_for_history(server, prompt_id, poll_seconds) + wall_seconds = time.perf_counter() - started + status = history.get("status", {}) + if status.get("status_str") == "error": + raise RuntimeError( + f"Sweep case {case['id']} failed: {json.dumps(status, indent=2)}" + ) + outputs = history.get("outputs", {}) + output_filename = api._text_output(outputs, linked[0]) + report = api._text_output(outputs, linked[2]) + metrics = parse_cwb_report(report) + metrics["wall_seconds"] = wall_seconds + result = { + **case, + "prompt_id": prompt_id, + "output_filename": output_filename, + "metrics": metrics, + "cwb_report": report, + } + if analysis_workflow is not None: + analysis_result = analysis_api.run_analysis( + deepcopy(analysis_workflow), server, poll_seconds + ) + metrics.update(parse_analysis_reports( + analysis_result["comparison_report"], + analysis_result["cwb_report"], + )) + result["analysis"] = analysis_result + return result + + +def _load_completed(path: Path) -> list[dict]: + if not path.exists(): + return [] + completed = [] + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + try: + completed.append(json.loads(line)) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"Invalid sweep JSONL at line {line_number}: {exc}" + ) from exc + return completed + + +def _append_result(path: Path, result: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(result, ensure_ascii=False) + "\n") + handle.flush() + os.fsync(handle.fileno()) + + +def _write_influence_csv(path: Path, results: list[dict]) -> None: + baseline = next((result for result in results if result["id"] == "baseline"), None) + if baseline is None: + return + fields = ["id", "factor", "value", *DELTA_METRICS] + delta_fields = [f"delta_{metric}" for metric in DELTA_METRICS] + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=[*fields, *delta_fields]) + writer.writeheader() + base_metrics = baseline["metrics"] + for result in results: + row = {key: result.get(key) for key in ("id", "factor", "value")} + for metric in DELTA_METRICS: + value = result["metrics"].get(metric) + base = base_metrics.get(metric) + row[metric] = value + row[f"delta_{metric}"] = ( + value - base if value is not None and base is not None else None + ) + writer.writerow(row) + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Run a resumable one-factor CWB LoRA parameter sweep." + ) + parser.add_argument("--workflow", type=Path, default=api.DEFAULT_WORKFLOW) + parser.add_argument("--server", default="http://127.0.0.1:8188") + parser.add_argument("--poll-seconds", type=float, default=2.0) + parser.add_argument("--output-dir", type=Path, default=Path("audit/cwb_sweep")) + parser.add_argument("--full-blocks", action="store_true") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--limit", type=int) + parser.add_argument( + "--case", + action="append", + dest="case_ids", + help="Run only the named case. Repeat to select multiple cases.", + ) + parser.add_argument( + "--analysis-workflow", + type=Path, + default=analysis_api.DEFAULT_WORKFLOW, + ) + parser.add_argument( + "--analysis-every", + type=int, + default=5, + help="Run candidate-versus-anchor analysis for baseline and every Nth case.", + ) + args = parser.parse_args() + + workflow = json.loads(args.workflow.read_text(encoding="utf-8")) + analysis_workflow = json.loads(args.analysis_workflow.read_text(encoding="utf-8")) + if args.full_blocks: + _, analysis_node = api._single_node( + analysis_workflow, + analysis_api.ANALYSIS_CLASS, + ) + analysis_node.setdefault("inputs", {})["exclude_patterns"] = "" + cases = build_sweep_cases(_baseline_config(workflow)) + if args.case_ids: + known_ids = {case["id"] for case in cases} + unknown = [case_id for case_id in args.case_ids if case_id not in known_ids] + if unknown: + raise RuntimeError("Unknown sweep case(s): " + ", ".join(unknown)) + selected = set(args.case_ids) + cases = [case for case in cases if case["id"] in selected] + if args.limit is not None: + cases = cases[:args.limit] + if args.dry_run: + for case in cases: + sys.stdout.write(f"{case['id']}: {case['factor']}={case['value']}\n") + sys.stdout.write(f"Cases: {len(cases)}\n") + return 0 + + results_path = args.output_dir / "runs.jsonl" + completed = _load_completed(results_path) + completed_ids = {result["id"] for result in completed} + server = args.server.rstrip("/") + if args.analysis_every < 1: + raise RuntimeError("--analysis-every must be at least 1") + for index, case in enumerate(cases, start=1): + if case["id"] in completed_ids: + sys.stdout.write(f"[{index}/{len(cases)}] skip {case['id']}\n") + continue + sys.stdout.write(f"[{index}/{len(cases)}] run {case['id']}\n") + result = _run_case( + workflow, + case, + server, + args.poll_seconds, + args.full_blocks, + analysis_workflow=( + analysis_workflow + if index == 1 or index % args.analysis_every == 0 + else None + ), + ) + _append_result(results_path, result) + completed.append(result) + completed_ids.add(case["id"]) + sys.stdout.write( + f" {result['metrics']['wall_seconds']:.1f}s, " + f"match={result['metrics']['alignment_match_fraction']:.4f}, " + f"effective={result['metrics']['effective_contributors_mean']:.4f}\n" + ) + _write_influence_csv(args.output_dir / "influence.csv", completed) + + _write_influence_csv(args.output_dir / "influence.csv", completed) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_consensus_merger.py b/tests/test_consensus_merger.py index f191d32..dc138e5 100644 --- a/tests/test_consensus_merger.py +++ b/tests/test_consensus_merger.py @@ -38,8 +38,7 @@ def _result_path(tmp_path, category, result): def _params(output_filename, **overrides): - params = { - "cwb_preset": "custom", + settings = { "consensus_type": "mean", "alignment_method": "index", "alignment_threshold": 0.4, @@ -52,6 +51,15 @@ def _params(output_filename, **overrides): "soft_comfort_bandpass": False, "position_weight": 0.0, "preserve_common_prefix": False, + } + for key in list(settings): + if key in overrides: + settings[key] = overrides.pop(key) + params = { + "cwb_preset": "broad_sim_medn_rn_softcb", + "cwb_config": _load_module("nodes.consensus_merger").build_custom_cwb_settings( + **settings + ), "mismatch_mode": "skip", "output_filename": output_filename, "save_dtype": "fp16", @@ -68,43 +76,68 @@ def _params(output_filename, **overrides): return params +def _settings(cwb, **overrides): + values = { + "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, + } + values.update(overrides) + return cwb.build_custom_cwb_settings(**values) + + def test_preset_resolution_matches_extended_contract(cwb): - assert all( - preset["global_scale"] == pytest.approx(1.0) - for preset in cwb.CWB_PRESETS.values() + default = cwb.resolve_cwb_settings( + "broad_sim_medn_rn_softcb", cwb.LORA_CWB_PRESETS ) + assert default.consensus_type == "median" + assert default.alignment_method == "similarity" + assert default.alignment_threshold == pytest.approx(0.0) + assert default.similarity_threshold == pytest.approx(0.0) + assert default.power_alpha == pytest.approx(2.0) + assert default.diversity_beta == pytest.approx(4.0) + assert default.rescale_norm is True + assert default.dynamic_similarity_contrast is False + assert default.soft_comfort_bandpass is True + assert default.position_weight == pytest.approx(0.05) - 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 + with pytest.raises(ValueError, match="Unsupported CWB preset"): + cwb.resolve_cwb_settings("baseline", cwb.LORA_CWB_PRESETS) - 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_preset_construction_rejects_omitted_fields(cwb): + with pytest.raises(TypeError): + cwb._preset( + consensus_type="median", + alignment_method="similarity", + alignment_threshold=0.0, + ) + + +def test_lora_alignment_presets_use_sweep_verified_thresholds(cwb): + assert cwb.LORA_CWB_PRESETS["broad_sim_medn_rn_softcb"][ + "alignment_threshold" + ] == pytest.approx(0.0) + assert cwb.LORA_CWB_PRESETS["moderate_sim_medn_rn_softcb"][ + "alignment_threshold" + ] == pytest.approx(0.0005) + assert cwb.LORA_CWB_PRESETS["conservative_sim_medn_rn_softcb"][ + "alignment_threshold" + ] == pytest.approx(0.0025) def test_cwb_math_index_prefix_and_scale(cwb): - settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + settings = _settings( + cwb, consensus_type="mean", alignment_method="index", global_scale=1.0, @@ -116,8 +149,8 @@ def test_cwb_math_index_prefix_and_scale(cwb): torch.tensor([[2.0, 0.0], [0.0, 2.0]]), ) - prefix_settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + prefix_settings = _settings( + cwb, consensus_type="mean", alignment_method="index", global_scale=2.0, @@ -133,8 +166,8 @@ def test_cwb_math_index_prefix_and_scale(cwb): def test_similarity_alignment_reorders_matching_rows(cwb): - settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + settings = _settings( + cwb, consensus_type="mean", alignment_method="similarity", alignment_threshold=0.5, @@ -152,8 +185,8 @@ def test_similarity_alignment_reorders_matching_rows(cwb): def test_fixed_coordinate_merge_never_reorders_model_rows(cwb): - settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + settings = _settings( + cwb, consensus_type="mean", alignment_method="similarity", alignment_threshold=0.5, @@ -177,8 +210,8 @@ def test_fixed_coordinate_merge_never_reorders_model_rows(cwb): def test_lora_similarity_alignment_pairs_a_rows_with_b_columns(cwb): - settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + settings = _settings( + cwb, consensus_type="mean", alignment_method="similarity", alignment_threshold=0.5, @@ -215,8 +248,8 @@ def test_lora_krea_shape_scores_only_rank_components(monkeypatch, cwb): "merge_consensus_group", lambda stacked, settings, **kwargs: stacked[0].clone(), ) - settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + settings = _settings( + cwb, alignment_method="similarity", alignment_threshold=0.5, ) @@ -230,9 +263,199 @@ def test_lora_krea_shape_scores_only_rank_components(monkeypatch, cwb): assert all(left[0] == rank and right[1] == rank for left, right in matrix_shapes) +def test_mixed_rank_similarity_searches_full_reference_without_padding( + monkeypatch, cwb +): + reference_down = torch.eye(3) + reference_up = torch.eye(3) + source_down = torch.tensor([[0.0, 0.0, 2.0]]) + source_up = torch.tensor([[0.0], [0.0], [4.0]]) + settings = _settings( + cwb, + consensus_type="mean", + alignment_method="similarity", + alignment_threshold=0.5, + rescale_norm=True, + ) + diagnostics = cwb.CWBDiagnostics() + matrix_shapes = [] + original_mm = cwb.torch.mm + + def record_mm(left, right): + matrix_shapes.append((tuple(left.shape), tuple(right.shape))) + return original_mm(left, right) + + monkeypatch.setattr(cwb.torch, "mm", record_mm) + merged_down, merged_up = cwb.merge_cwb_lora_pairs( + [reference_down, source_down], + [reference_up, source_up], + settings, + reference_index=0, + diagnostics=diagnostics, + ) + + assert matrix_shapes == [((3, 3), (3, 1)), ((3, 3), (3, 1))] + assert diagnostics.alignment_matches == 1 + assert diagnostics.anchor_only_groups == 2 + torch.testing.assert_close(merged_down[:2], reference_down[:2]) + torch.testing.assert_close(merged_up[:, :2], reference_up[:, :2]) + torch.testing.assert_close(merged_down[2], torch.tensor([0.0, 0.0, 1.5])) + torch.testing.assert_close(merged_up[:, 2], torch.tensor([0.0, 0.0, 2.5])) + + +def test_mixed_rank_index_does_not_rescale_absent_components(cwb): + reference_down = torch.eye(3) + reference_up = torch.eye(3) + source_down = torch.tensor([[2.0, 0.0, 0.0]]) + source_up = torch.tensor([[2.0], [0.0], [0.0]]) + settings = _settings( + cwb, + consensus_type="mean", + alignment_method="index", + rescale_norm=True, + ) + + merged_down, merged_up = cwb.merge_cwb_lora_pairs( + [reference_down, source_down], + [reference_up, source_up], + settings, + reference_index=0, + ) + + torch.testing.assert_close(merged_down[1:], reference_down[1:]) + torch.testing.assert_close(merged_up[:, 1:], reference_up[:, 1:]) + assert torch.linalg.vector_norm(merged_down[0]).item() == pytest.approx(1.5) + assert torch.linalg.vector_norm(merged_up[:, 0]).item() == pytest.approx(1.5) + + +def test_genuine_zero_component_is_not_treated_as_structural_padding(cwb): + reference_down = torch.eye(2) + reference_up = torch.eye(2) + source_down = torch.zeros((1, 2)) + source_up = torch.zeros((2, 1)) + settings = _settings( + cwb, + alignment_method="similarity", + alignment_threshold=0.0, + ) + diagnostics = cwb.CWBDiagnostics() + + cwb.merge_cwb_lora_pairs( + [reference_down, source_down], + [reference_up, source_up], + settings, + reference_index=0, + diagnostics=diagnostics, + ) + + assert diagnostics.alignment_matches == 1 + assert diagnostics.anchor_only_groups == 1 + + +def test_explicit_mismatch_zero_applies_to_every_reference_component(cwb): + down = torch.eye(2) + up = torch.eye(2) + settings = _settings( + cwb, + consensus_type="mean", + alignment_method="index", + rescale_norm=True, + ) + diagnostics = cwb.CWBDiagnostics() + + merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu( + [cwb.LoRAPairSource(0, down, up), cwb.LoRAPairSource.zero(1)], + settings, + "cpu", + torch.float32, + diagnostics=diagnostics, + ) + + assert rank == 2 + torch.testing.assert_close(merged_down, down * 0.5) + torch.testing.assert_close(merged_up, up * 0.5) + assert diagnostics.explicit_zero_contributors == 1 + assert diagnostics.structural_rank_slots_excluded == 0 + + +def test_mixed_rank_diagnostics_count_only_genuine_components(cwb): + settings = _settings(cwb, alignment_method="index") + diagnostics = cwb.CWBDiagnostics() + sources = [ + cwb.LoRAPairSource(0, torch.ones((4, 2)), torch.ones((2, 4))), + cwb.LoRAPairSource(1, torch.ones((2, 2)), torch.ones((2, 2))), + cwb.LoRAPairSource(2, torch.ones((3, 2)), torch.ones((2, 3))), + ] + + _, _, rank = cwb._merge_lora_pair_to_cpu( + sources, + settings, + "cpu", + torch.float32, + diagnostics=diagnostics, + ) + + assert rank == 4 + assert diagnostics.lora_groups == 1 + assert diagnostics.mixed_rank_lora_groups == 1 + assert diagnostics.genuine_lora_components == 9 + assert diagnostics.structural_rank_slots_excluded == 3 + report = diagnostics.render("custom", 3, True) + assert "Mixed-rank LoRA groups: 1" in report + assert "Genuine LoRA components: 9" in report + assert "Structural rank slots excluded: 3" in report + assert "CWB LORA DIMENSION GROUPS" in report + assert "Input ranks: (4, 2, 3)" in report + assert "Matched source components: 5/5 (100.00%)" in report + + +def test_lora_report_groups_dimensions_and_summarizes_rectangular_similarity(cwb): + settings = _settings( + cwb, + consensus_type="mean", + alignment_method="similarity", + alignment_threshold=0.5, + ) + diagnostics = cwb.CWBDiagnostics() + sources = [ + cwb.LoRAPairSource(0, torch.eye(3), torch.eye(3)), + cwb.LoRAPairSource( + 1, + torch.tensor([[0.0, 0.0, 2.0]]), + torch.tensor([[0.0], [0.0], [2.0]]), + ), + ] + + for layer in ("blocks.1.attn", "blocks.2.attn"): + cwb._merge_lora_pair_to_cpu( + sources, + settings, + "cpu", + torch.float32, + operation_label=layer, + diagnostics=diagnostics, + ) + + report = diagnostics.render("custom", 2, True) + assert "Group 1: 2 layer(s)" in report + assert "Alignment: similarity; matrices: 3x1" in report + assert "Matched source components: 2/2 (100.00%)" in report + assert "Reference-only components: 4/6" in report + assert "Structural rank slots excluded: 4" in report + assert "Candidate similarity count: 6" in report + assert "Matched similarity p05/median/p95: 1 / 1 / 1" in report + assert "blocks.1.attn" in report + assert "blocks.2.attn" in report + assert all( + not isinstance(value, torch.Tensor) + for layer_report in diagnostics.lora_layer_reports + for value in vars(layer_report).values() + ) + + def test_lora_alpha_and_global_scale_apply_once_to_pair(cwb): - settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + settings = _settings( + cwb, consensus_type="mean", alignment_method="index", global_scale=3.0, @@ -240,7 +463,10 @@ def test_lora_alpha_and_global_scale_apply_once_to_pair(cwb): down = torch.ones((1, 1)) up = torch.ones((1, 1)) merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu( - [(0, down, up, 2.0), (1, down, up, 2.0)], + [ + cwb.LoRAPairSource(0, down, up, 2.0), + cwb.LoRAPairSource(1, down, up, 2.0), + ], settings, "cpu", torch.float32, @@ -251,8 +477,8 @@ def test_lora_alpha_and_global_scale_apply_once_to_pair(cwb): def test_lora_pair_alignment_supports_convolution_factors(cwb): - settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + settings = _settings( + cwb, consensus_type="mean", alignment_method="index", ) @@ -266,8 +492,8 @@ def test_lora_pair_alignment_supports_convolution_factors(cwb): def test_lora_cuda_oom_retries_current_pair_on_cpu(monkeypatch, caplog, cwb): - settings = cwb.resolve_cwb_settings( - cwb_preset="custom", + settings = _settings( + cwb, consensus_type="mean", alignment_method="index", ) @@ -282,29 +508,64 @@ def test_lora_cuda_oom_retries_current_pair_on_cpu(monkeypatch, caplog, cwb): monkeypatch.setattr(cwb, "_to_compute", fail_cuda) monkeypatch.setattr(cwb, "_release_failed_cuda_operation", lambda: None) - down = torch.ones((1, 2)) - up = torch.ones((2, 1)) + reference_down = torch.eye(2) + reference_up = torch.eye(2) + source_down = torch.ones((1, 2)) + source_up = torch.ones((2, 1)) + diagnostics = cwb.CWBDiagnostics() with caplog.at_level("WARNING"): merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu( - [(0, down, up, 1.0), (1, down, up, 1.0)], + [ + cwb.LoRAPairSource(0, reference_down, reference_up), + cwb.LoRAPairSource(1, source_down, source_up), + ], settings, "cuda", torch.float32, operation_label="diffusion_model.foo", + diagnostics=diagnostics, ) - assert rank == 1 - torch.testing.assert_close(merged_down, down) - torch.testing.assert_close(merged_up, up) + assert rank == 2 + assert merged_down.shape == reference_down.shape + assert merged_up.shape == reference_up.shape assert attempted_devices[0] == "cuda" assert "cpu" in attempted_devices assert "retrying this layer on CPU" in caplog.text assert "diffusion_model.foo" in caplog.text + assert diagnostics.cpu_fallbacks == 1 + assert diagnostics.lora_groups == 1 + assert diagnostics.structural_rank_slots_excluded == 1 + + +def test_mixed_rank_common_prefix_is_limited_to_smallest_rank(cwb): + reference_down = torch.eye(3) + reference_up = torch.eye(3) + source_down = torch.tensor([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0]]) + source_up = torch.tensor([[1.0, 0.0], [0.0, 2.0], [0.0, 0.0]]) + settings = _settings( + cwb, + consensus_type="mean", + alignment_method="index", + preserve_common_prefix=True, + ) + + merged_down, merged_up = cwb.merge_cwb_lora_pairs( + [reference_down, source_down], + [reference_up, source_up], + settings, + reference_index=0, + ) + + torch.testing.assert_close(merged_down[0], reference_down[0]) + torch.testing.assert_close(merged_up[:, 0], reference_up[:, 0]) + torch.testing.assert_close(merged_down[2], reference_down[2]) + torch.testing.assert_close(merged_up[:, 2], reference_up[:, 2]) def test_norm_rescale_and_dsc_bandpass_are_finite(cwb): - rescale = cwb.resolve_cwb_settings( - cwb_preset="custom", + rescale = _settings( + cwb, consensus_type="mean", rescale_norm=True, ) @@ -313,8 +574,8 @@ def test_norm_rescale_and_dsc_bandpass_are_finite(cwb): ) assert torch.linalg.vector_norm(merged).item() == pytest.approx(2.0) - dsc = cwb.resolve_cwb_settings( - cwb_preset="custom", + dsc = _settings( + cwb, consensus_type="median", diversity_beta=1.5, dynamic_similarity_contrast=True, @@ -326,14 +587,42 @@ def test_norm_rescale_and_dsc_bandpass_are_finite(cwb): 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 +def test_all_cwb_schemas_and_compact_control_contract(cwb): + assert len(cwb.CWB_MERGER_NODES) == 12 + schemas = [node.define_schema() for node in cwb.CWB_MERGER_NODES] + assert len({schema.node_id for schema in schemas}) == 12 + advanced = { + "consensus_type", "alignment_method", "alignment_threshold", + "similarity_threshold", "power_alpha", "diversity_beta", + "rescale_norm", "global_scale", "dynamic_similarity_contrast", + "soft_comfort_bandpass", "position_weight", "preserve_common_prefix", + } + config_schema = cwb.CWBCustomConfiguration.define_schema() + assert {value.id for value in config_schema.inputs} == advanced + for node in cwb.CWB_MERGER_NODES: schema = node.define_schema() ids = [value.id for value in schema.inputs] assert schema.description assert all(value.tooltip for value in schema.inputs) + if node is cwb.CWBCustomConfiguration: + continue + assert [output.display_name for output in schema.outputs] == [ + "output_filename", "documentation", "cwb_report" + ] + assert schema.outputs[0].get_io_type() == "*" + assert schema.outputs[1].get_io_type() == "STRING" + assert schema.outputs[2].get_io_type() == "STRING" + assert advanced.isdisjoint(ids) + assert "cwb_config" in ids + if node is cwb.CWBLoRAMultiMerger: + assert ids[:4] == ["execution_mode", "lora_count", "lora_1", "lora_2"] + assert schema.inputs[1].default == "2" + assert ids[-2:] == [ + "include_1d_diffs", + "counterfactual_weight_sweep", + ] + continue assert ids[:3] == ["execution_mode", "model_a", "model_b"] if node.INPUT_COUNT == 3: assert ids[3] == "model_c" @@ -344,6 +633,276 @@ def test_all_ten_schemas_and_lora_switch_contract(cwb): assert "include_1d_diffs" not in ids +def test_use_case_preset_names_match_complete_settings(cwb): + fields = set(cwb.CWBSettings.__dataclass_fields__) + for registry in ( + cwb.DENSE_CWB_PRESETS, + cwb.EMBEDDING_CWB_PRESETS, + cwb.LORA_CWB_PRESETS, + ): + assert 6 <= len(registry) <= 8 + for name, values in registry.items(): + assert set(values) == fields + assert ("_rn" in name) is values["rescale_norm"] + assert ("_dsc" in name) is values["dynamic_similarity_contrast"] + assert ("_softcb" in name) is values["soft_comfort_bandpass"] + assert ("_pcp" in name) is values["preserve_common_prefix"] + if registry is not cwb.DENSE_CWB_PRESETS: + expected_alignment = "similarity" if "_sim_" in name else "index" + assert values["alignment_method"] == expected_alignment + expected_consensus = "median" if "_medn" in name else "mean" + assert values["consensus_type"] == expected_consensus + + +def test_connected_custom_config_completely_overrides_preset(cwb): + custom = _settings( + cwb, + consensus_type="median", + alignment_method="similarity", + alignment_threshold=0.17, + similarity_threshold=-0.25, + diversity_beta=7.0, + dynamic_similarity_contrast=True, + ) + resolved = cwb.resolve_cwb_settings( + "direct_idx_medn_rn_softcb", cwb.LORA_CWB_PRESETS, custom + ) + assert resolved is custom + + +def test_diagnostics_distinguish_alignment_and_weight_fallbacks(cwb): + alignment_diagnostics = cwb.CWBDiagnostics() + settings = _settings( + cwb, alignment_method="similarity", alignment_threshold=0.9 + ) + cwb.merge_cwb_tensors( + [torch.eye(2), -torch.eye(2)], + settings, + diagnostics=alignment_diagnostics, + ) + assert alignment_diagnostics.alignment_matches == 0 + assert alignment_diagnostics.anchor_only_groups == 2 + + weighting_diagnostics = cwb.CWBDiagnostics() + cwb.merge_consensus_group( + torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + _settings(cwb, similarity_threshold=0.9), + diagnostics=weighting_diagnostics, + ) + assert weighting_diagnostics.all_rejected_fallbacks == 1 + assert weighting_diagnostics.weighting_contributors == 2 + assert weighting_diagnostics.accepted_weighting_contributors == 0 + assert weighting_diagnostics.dominant_weight_min == pytest.approx(0.5) + assert weighting_diagnostics.dominant_weight_max == pytest.approx(0.5) + assert weighting_diagnostics.effective_contributor_min == pytest.approx(2.0) + assert weighting_diagnostics.effective_contributor_max == pytest.approx(2.0) + + +def test_weight_diagnostics_measure_parameter_influence(cwb): + diagnostics = cwb.CWBDiagnostics() + cwb.merge_consensus_group( + torch.tensor([[1.0, 0.0], [0.8, 0.2], [-1.0, 0.0]]), + _settings( + cwb, + similarity_threshold=0.0, + power_alpha=2.0, + diversity_beta=0.0, + ), + diagnostics=diagnostics, + ) + + assert diagnostics.weighting_groups == 1 + assert diagnostics.weighting_contributors == 3 + assert diagnostics.accepted_weighting_contributors == 2 + assert diagnostics.dominant_weight_max > 0.5 + assert 1.0 < diagnostics.effective_contributor_min < 2.0 + report = diagnostics.render("test", 3, custom=False) + assert "Accepted weighting contributors: 2/3" in report + assert "Dominant normalized weight min/mean/max:" in report + assert "Effective contributors min/mean/max:" in report + + +def test_counterfactual_weight_sweep_reuses_similarity_vectors(cwb): + sweep = cwb.CWBWeightSweepDiagnostics() + similarities = torch.tensor([0.25, 0.5, 0.75]) + sweep._record_similarities("median", similarities) + + baseline = sweep.stats[("median", 0.0, 2.0, 0.0, False, False)] + expected = similarities.square() + expected /= expected.sum() + assert baseline.groups == 1 + assert baseline.accepted == 3 + assert baseline.dominant_sum == pytest.approx(float(expected.max())) + assert baseline.effective_sum == pytest.approx( + float(1.0 / expected.square().sum()) + ) + assert len(sweep.stats) == 180 + assert all(not isinstance(value, torch.Tensor) for value in sweep.__dict__.values()) + + +def test_default_lora_preset_changes_perturbed_factors(cwb): + settings = cwb.resolve_cwb_settings( + "broad_sim_medn_rn_softcb", cwb.LORA_CWB_PRESETS + ) + reference_down = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + reference_up = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + source_down = torch.tensor([[1.0, 0.1], [0.1, 1.0]]) + source_up = torch.tensor([[1.1, 0.1], [0.1, 0.9]]) + diagnostics = cwb.CWBDiagnostics() + merged_down, merged_up = cwb.merge_cwb_lora_pairs( + [reference_down, source_down], + [reference_up, source_up], + settings, + reference_index=0, + diagnostics=diagnostics, + ) + assert diagnostics.alignment_matches == 2 + assert diagnostics.anchor_only_groups == 0 + assert not torch.equal(merged_down, reference_down) + assert not torch.equal(merged_up, reference_up) + + +def test_eight_input_lora_merge_normalizes_before_variable_rank_alignment( + monkeypatch, tmp_path, cwb +): + paths = {} + events = [] + original_normalize = cwb.normalize_lora_pair + original_merge = cwb.merge_cwb_lora_pairs + original_mark_processed = cwb.MemoryEfficientSafeOpen.mark_processed + received_ranks = [] + processed = [] + for index in range(8): + rank = 1 if index < 4 else 2 + name = f"lora_{index}" + path = tmp_path / f"{name}.safetensors" + save_file({ + "diffusion_model.foo.lora_A.weight": torch.full( + (rank, 2), 1.0 + index * 0.01 + ), + "diffusion_model.foo.lora_B.weight": torch.full( + (2, rank), 1.0 + index * 0.02 + ), + "diffusion_model.foo.alpha": torch.tensor(float(rank) * 0.5), + }, str(path)) + paths[name] = str(path) + + def record_normalize(*args, **kwargs): + events.append("normalize") + return original_normalize(*args, **kwargs) + + def record_merge(downs, ups, *args, **kwargs): + events.append("merge") + received_ranks.extend(down.shape[0] for down in downs) + for down, up in zip(downs, ups): + assert down.shape[0] == up.shape[1] + return original_merge(downs, ups, *args, **kwargs) + + def record_processed(loader, key): + processed.append(key) + return original_mark_processed(loader, key) + + monkeypatch.setattr(cwb, "normalize_lora_pair", record_normalize) + monkeypatch.setattr(cwb, "merge_cwb_lora_pairs", record_merge) + monkeypatch.setattr( + cwb.MemoryEfficientSafeOpen, "mark_processed", record_processed + ) + _patch_io(monkeypatch, cwb, tmp_path, paths) + result = cwb.ConsensusMergerLogic.execute( + list(paths), "loras", _params("eight_input"), lora_mode=True + ) + tensors = load_file(str(_result_path(tmp_path, "loras", result))) + assert tensors["diffusion_model.foo.lora_A.weight"].shape[0] == 2 + assert "diffusion_model.foo.alpha" not in tensors + assert events.count("normalize") == 8 + assert events.index("merge") > max( + index for index, event in enumerate(events) if event == "normalize" + ) + assert received_ranks == [1, 1, 1, 1, 2, 2, 2, 2] + assert len(processed) == 24 + assert processed.count("diffusion_model.foo.lora_A.weight") == 8 + assert processed.count("diffusion_model.foo.lora_B.weight") == 8 + assert processed.count("diffusion_model.foo.alpha") == 8 + + +def test_lora_mismatch_zeros_remains_explicit_across_full_output_rank( + monkeypatch, tmp_path, cwb +): + a = tmp_path / "mismatch_zero_a.safetensors" + b = tmp_path / "mismatch_zero_b.safetensors" + save_file({ + "diffusion_model.foo.lora_A.weight": torch.eye(2), + "diffusion_model.foo.lora_B.weight": torch.eye(2), + }, str(a)) + save_file({"unrelated": torch.tensor([1.0])}, str(b)) + _patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)}) + + result = cwb.ConsensusMergerLogic.execute( + ["a", "b"], + "loras", + _params( + "mismatch_zero", + mismatch_mode="zeros", + consensus_type="mean", + alignment_method="index", + rescale_norm=True, + save_dtype="fp32", + ), + lora_mode=True, + ) + tensors = load_file(str(_result_path(tmp_path, "loras", result))) + + torch.testing.assert_close( + tensors["diffusion_model.foo.lora_A.weight"], torch.eye(2) * 0.5 + ) + torch.testing.assert_close( + tensors["diffusion_model.foo.lora_B.weight"], torch.eye(2) * 0.5 + ) + + +def test_lora_multi_wrapper_validates_count_and_forwards_equal_prior_names( + monkeypatch, cwb +): + kwargs = { + "execution_mode": "MERGE", + "lora_count": "3", + "lora_1": "a.safetensors", + "lora_2": "b.safetensors", + "lora_3": "None", + "lora_4": "None", + "lora_5": "None", + "lora_6": "None", + "lora_7": "None", + "lora_8": "None", + "cwb_preset": "broad_sim_medn_rn_softcb", + "cwb_config": None, + } + monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS") + with pytest.raises(ValueError, match="unselected input.*3"): + cwb.CWBLoRAMultiMerger.execute(**kwargs) + + 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) + kwargs["lora_3"] = "c.safetensors" + result = cwb.CWBLoRAMultiMerger.execute(**kwargs) + assert captured["names"] == [ + "a.safetensors", "b.safetensors", "c.safetensors" + ] + assert captured["model_type"] == "loras" + assert captured["options"] == {"lora_mode": True} + assert result.result == ("merged.safetensors", "DOCS", "REPORT") + + def test_streaming_dtype_nonfloat_and_secondary_preservation( monkeypatch, tmp_path, cwb ): @@ -436,6 +995,25 @@ def test_embedding_union_uses_longest_first_dimension(monkeypatch, tmp_path, cwb torch.testing.assert_close(tensors["secondary_only"], torch.tensor([[7.0, 0.0]])) +def test_output_name_preserves_category_relative_subdirectory( + monkeypatch, tmp_path, cwb +): + a = tmp_path / "relative_a.safetensors" + b = tmp_path / "relative_b.safetensors" + save_file({"layer": torch.tensor([[1.0]])}, str(a)) + save_file({"layer": torch.tensor([[1.0]])}, str(b)) + _patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)}) + + result = cwb.ConsensusMergerLogic.execute( + ["a", "b"], + "loras", + _params("groupfolder/modelname"), + ) + + assert result == "groupfolder/modelname.safetensors" + assert (tmp_path / "loras" / "groupfolder" / "modelname.safetensors").is_file() + + def test_three_input_streaming_merge(monkeypatch, tmp_path, cwb): paths = {} for name, value in (("a", 1.0), ("b", 3.0), ("c", 5.0)): diff --git a/tests/test_cwb_api_runner.py b/tests/test_cwb_api_runner.py new file mode 100644 index 0000000..433eb96 --- /dev/null +++ b/tests/test_cwb_api_runner.py @@ -0,0 +1,141 @@ +import importlib.util +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _load_runner(): + spec = importlib.util.spec_from_file_location( + "cwb_api_runner_test", REPO_ROOT / "run_cwb_api_workflow.py" + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def runner(): + return _load_runner() + + +@pytest.fixture +def workflow(): + return { + "10": { + "class_type": "CWBCustomConfiguration", + "inputs": {"diversity_beta": 10.0}, + }, + "11": { + "class_type": "CWBLoRAMultiMerger", + "inputs": { + "cwb_config": ["10", 0], + "cwb_preset": "default", + "output_filename": "original", + }, + }, + "14": {"class_type": "PreviewAny", "inputs": {"source": ["11", 0]}}, + "15": {"class_type": "PreviewAny", "inputs": {"source": ["11", 2]}}, + } + + +def test_custom_export_remains_connected_by_default(runner, workflow): + merger_id, linked = runner._prepare_workflow(workflow, None, None) + assert merger_id == "11" + assert workflow["11"]["inputs"]["cwb_config"] == ["10", 0] + assert workflow["11"]["inputs"]["exclude_patterns"] == ( + runner.CALIBRATION_EXCLUDE_PATTERN + ) + assert linked == {0: "14", 2: "15"} + + +def test_preset_run_disconnects_custom_config_and_overrides_filename( + runner, workflow +): + runner._prepare_workflow(workflow, "candidate_sim_mean", "candidate/output") + inputs = workflow["11"]["inputs"] + assert "cwb_config" not in inputs + assert inputs["cwb_preset"] == "candidate_sim_mean" + assert inputs["output_filename"] == "candidate/output" + + +def test_full_block_verification_does_not_add_calibration_exclusions( + runner, workflow +): + workflow["11"]["inputs"]["exclude_patterns"] = "existing" + runner._prepare_workflow(workflow, None, None, full_blocks=True) + assert workflow["11"]["inputs"]["exclude_patterns"] == "existing" + + +@pytest.mark.parametrize( + ("block", "excluded"), + [(block, block not in {0, 3, 14, 24, 27}) for block in range(28)], +) +def test_calibration_block_ranges_are_exact(runner, block, excluded): + import re + + key = f"diffusion_model.blocks.{block}.attn.wq.lora_A.weight" + patterns = runner.CALIBRATION_EXCLUDE_PATTERN.splitlines() + assert any(re.search(pattern, key) for pattern in patterns) is excluded + + +@pytest.mark.parametrize( + ("key", "excluded"), + [ + ("diffusion_model.txtfusion.layerwise_blocks.1.attn.lora_A.weight", True), + ("diffusion_model.txtfusion.refiner_blocks.0.attn.lora_A.weight", True), + ("diffusion_model.txtfusion.layerwise_blocks.0.attn.lora_A.weight", False), + ], +) +def test_calibration_txtfusion_exclusions_are_exact(runner, key, excluded): + import re + + patterns = runner.CALIBRATION_EXCLUDE_PATTERN.splitlines() + assert any(re.search(pattern, key) for pattern in patterns) is excluded + + +def test_requires_filename_and_report_capture(runner, workflow): + del workflow["15"] + with pytest.raises(RuntimeError, match="output 2"): + runner._prepare_workflow(workflow, None, None) + + +def test_parses_counterfactual_weight_sweep(runner): + report = """CWB COUNTERFACTUAL WEIGHT SWEEP +median | 0.1 | 2 | 4 | true | false | 12/18 | 3 | 0.6/1 | 2.4/1 +""" + rows = runner.parse_counterfactual_weight_sweep(report) + assert rows == [{ + "consensus_type": "median", + "similarity_threshold": 0.1, + "power_alpha": 2.0, + "diversity_beta": 4.0, + "dynamic_similarity_contrast": True, + "soft_comfort_bandpass": False, + "accepted_contributors": 12, + "contributors": 18, + "fallbacks": 3, + "dominant_weight_mean": 0.6, + "dominant_weight_max": 1.0, + "effective_contributors_mean": 2.4, + "effective_contributors_min": 1.0, + }] + + +def test_prompt_payload_contains_only_api_workflow(runner, workflow): + payload = runner._prompt_payload(workflow) + assert payload == {"prompt": workflow} + + +@pytest.mark.parametrize( + ("output", "expected"), + [ + ({"text": ["value"]}, "value"), + ({"text": "value"}, "value"), + ({"ui": {"text": ["value"]}}, "value"), + ], +) +def test_extracts_supported_preview_history_shapes(runner, output, expected): + assert runner._text_output({"14": output}, "14") == expected diff --git a/tests/test_cwb_preset_sweep.py b/tests/test_cwb_preset_sweep.py new file mode 100644 index 0000000..2823d57 --- /dev/null +++ b/tests/test_cwb_preset_sweep.py @@ -0,0 +1,120 @@ +import importlib.util +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _load_sweep(): + spec = importlib.util.spec_from_file_location( + "cwb_preset_sweep_test", REPO_ROOT / "run_cwb_preset_sweep.py" + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def sweep(): + return _load_sweep() + + +@pytest.fixture +def baseline(): + return { + "consensus_type": "median", + "alignment_method": "similarity", + "alignment_threshold": 0.0, + "similarity_threshold": 0.0, + "power_alpha": 2.0, + "diversity_beta": 10.0, + "rescale_norm": True, + "global_scale": 1.0, + "dynamic_similarity_contrast": False, + "soft_comfort_bandpass": True, + "position_weight": 0.05, + "preserve_common_prefix": False, + } + + +def test_cases_are_complete_and_change_exactly_one_factor(sweep, baseline): + cases = sweep.build_sweep_cases(baseline) + assert cases[0]["id"] == "baseline" + assert len({case["id"] for case in cases}) == len(cases) + for case in cases: + assert tuple(case["config"]) == sweep.CONFIG_KEYS + changed = [ + key for key in sweep.CONFIG_KEYS + if case["config"][key] != baseline[key] + ] + assert changed == ([] if case["id"] == "baseline" else [case["factor"]]) + + +def test_sweep_exposes_expected_full_block_finalist_ids(sweep, baseline): + ids = {case["id"] for case in sweep.build_sweep_cases(baseline)} + assert { + "baseline", + "alignment_method_index", + "alignment_threshold_0p0005", + "alignment_threshold_0p0025", + "consensus_type_mean", + "diversity_beta_4p0", + "rescale_norm_false", + "dynamic_similarity_contrast_true", + "soft_comfort_bandpass_false", + } <= ids + + +def test_report_parser_extracts_influence_metrics(sweep): + report = """CWB MERGE SUMMARY +Alignment candidates: 100 +Alignment matches: 75 +Anchor-only groups: 5 +Weighted vector groups: 20 +All-rejected equal fallbacks: 1 +Zero-weight equal fallbacks: 2 +Accepted weighting contributors: 50/80 +Consensus similarity min/mean/max: -0.5 / 0.25 / 1 +Dominant normalized weight min/mean/max: 0.2 / 0.4 / 0.8 +Effective contributors min/mean/max: 1.2 / 2.5 / 5 +CUDA OOM CPU fallbacks: 0 +LoRA groups: 10 +Rank-one delta norm ratio count/min/mean/max: 3 / 0.7 / 1.1 / 1.4 +Rank-one delta norm ratio count/min/mean/max: 1 / 0.8 / 0.9 / 1.2 +""" + metrics = sweep.parse_cwb_report(report) + assert metrics["alignment_match_fraction"] == pytest.approx(0.75) + assert metrics["accepted_weight_fraction"] == pytest.approx(0.625) + assert metrics["dominant_weight_mean"] == pytest.approx(0.4) + assert metrics["effective_contributors_mean"] == pytest.approx(2.5) + assert metrics["norm_ratio_min"] == pytest.approx(0.7) + assert metrics["norm_ratio_count"] == 4 + assert metrics["norm_ratio_mean"] == pytest.approx(1.05) + assert metrics["norm_ratio_max"] == pytest.approx(1.4) + + +def test_report_parser_rejects_pre_diagnostic_server(sweep): + with pytest.raises(RuntimeError, match="Restart ComfyUI"): + sweep.parse_cwb_report("CWB MERGE SUMMARY\nAlignment matches: 1") + + +def test_analysis_parser_extracts_global_metrics(sweep): + comparison = """MODEL COMPARISON SUMMARY +PARAMETER-WEIGHTED GLOBAL METRICS +relative_l2: 1.04 +cosine: 0.46 +exact: 0.01 +norm_ratio: 0.89 +EQUAL-LAYER AVERAGES +""" + cwb_report = """GLOBAL CWB-DERIVED DIAGNOSTICS +Pairwise row cosine mean/min/max: 0.4 / -0.1 / 1 +Mean-consensus affinity A/B: 0.2 / 0.3 +Median-consensus affinity A/B: 0.25 / 0.35 +""" + metrics = sweep.parse_analysis_reports(comparison, cwb_report) + assert metrics["analysis_relative_l2"] == pytest.approx(1.04) + assert metrics["analysis_pair_cosine"] == pytest.approx(0.4) + assert metrics["analysis_median_affinity_b"] == pytest.approx(0.35)