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

feat(cwb): add LoRA presets and diagnostics

Browse Source
This commit is contained in:
silveroxides
2026-08-14 20:53:42 +02:00
parent ece52e651c
commit c25d6c4990
8 changed files with 2798 additions and 540 deletions
Download Patch File Download Diff File Expand all files Collapse all files
audit/uel_multi_input_loading_audit.md
-136
View File
@@ -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.
docs/consensus_mergers.md
+107 -40
View File
@@ -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.<layer>.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
nodes/consensus_merger.py
+1114 -303
View File
File diff suppressed because it is too large Load Diff
run_cwb_api_workflow.py
+259
View File
@@ -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"^(?P<consensus_type>mean|median) \| "
r"(?P<similarity_threshold>[-+0-9.e]+) \| "
r"(?P<power_alpha>[-+0-9.e]+) \| (?P<diversity_beta>[-+0-9.e]+) \| "
r"(?P<dsc>true|false) \| (?P<softcb>true|false) \| "
r"(?P<accepted>\d+)/(?P<contributors>\d+) \| (?P<fallbacks>\d+) \| "
r"(?P<dominant_mean>[-+0-9.e]+)/(?P<dominant_max>[-+0-9.e]+) \| "
r"(?P<effective_mean>[-+0-9.e]+)/(?P<effective_min>[-+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())
run_cwb_preset_sweep.py
+418
View File
@@ -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())
tests/test_consensus_merger.py
+639 -61
View File
@@ -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)):
tests/test_cwb_api_runner.py
+141
View File
@@ -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
tests/test_cwb_preset_sweep.py
+120
View File
@@ -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)