From e8eb2a08681eb87d018a17ae3e7468136ca566ea Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sun, 13 Sep 2026 14:56:17 +0200 Subject: [PATCH] feat: Analyze LoRA effects on a base model --- __init__.py | 2 + docs/lora_model_analysis.md | 30 ++++ nodes/lora_model_analysis.py | 241 ++++++++++++++++++++++++++++++ nodes/lora_resize.py | 27 ++-- tests/test_groups.toml | 4 + tests/test_lora_model_analysis.py | 234 +++++++++++++++++++++++++++++ 6 files changed, 525 insertions(+), 13 deletions(-) create mode 100644 docs/lora_model_analysis.md create mode 100644 nodes/lora_model_analysis.py create mode 100644 tests/test_lora_model_analysis.py diff --git a/__init__.py b/__init__.py index 9428a7e..f634ba9 100644 --- a/__init__.py +++ b/__init__.py @@ -62,6 +62,7 @@ from .nodes.model_info_nodes import ( from .nodes.lora_rename import AnimaLoraRename from .nodes.dtype_conversion import DiffusionModelDtypeConversion from .nodes.layer_parameters import LayerParameterConfiguration +from .nodes.lora_model_analysis import LoRAOnModelAnalysis class ModelUtilsExtension(ComfyExtension): @override @@ -91,6 +92,7 @@ class ModelUtilsExtension(ComfyExtension): *LODESTONE_MERGER_NODES, # Two-Model Analysis *MODEL_ANALYSIS_NODES, + LoRAOnModelAnalysis, # LoRA Extraction LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractKnee, LoRAExtractFrobenius, diff --git a/docs/lora_model_analysis.md b/docs/lora_model_analysis.md new file mode 100644 index 0000000..2ec05c5 --- /dev/null +++ b/docs/lora_model_analysis.md @@ -0,0 +1,30 @@ +# LoRA-on-model analysis + +**Analyze LoRA Effect on Diffusion Model** takes a LoRA as `model_a` and a base +diffusion model as `model_b`. It applies the LoRA transiently and compares the +original base weights against the patched weights. The measured update is +`patched − original`; no merged checkpoint, delta file, or inference is produced. + +The comparison and CWB reports use original base as A and patched base as B. +The CSV outputs contain per-tensor difference metrics and row cosine summaries. +Unchanged floating base tensors remain in the analysis, so global metrics reflect +the impact on the selected base model, not just the layers touched by the LoRA. +Application and statistics use FP32, matching the existing analysis arithmetic; +this does not simulate rounding to a separately saved model dtype. + +LoRA strength defaults to 1. Alpha/rank normalization, DoRA, LoCon, and reshape +handling use the existing application path and ComfyUI adapter. Direct deltas, +including bias deltas, are applied; explicit `set_weight` patches replace weights +regardless of strength, matching the existing application behavior. + +Exclude Patterns and Include Mode filter **base tensor keys** before loading. +Unmapped LoRA groups and unused source tensors are listed separately. Unsupported +low-bit work units and non-floating base tensors are reported and omitted from +numeric metrics; they are not decoded or represented as successfully analyzed. +Application errors fail the operation rather than report a failed patch as zero +impact. A CUDA out-of-memory error retries the complete affected unit on CPU. + +Loading uses the shared UEL work-unit streamer: one base tensor and its required +adapter tensors, batch size 1 and prefetch depth 1. Sources are marked processed +and streams closed on success or failure. Only scalar statistics and bounded top +differences survive each unit. No whole model or LoRA tensor dictionary is built. diff --git a/nodes/lora_model_analysis.py b/nodes/lora_model_analysis.py new file mode 100644 index 0000000..53b3ce2 --- /dev/null +++ b/nodes/lora_model_analysis.py @@ -0,0 +1,241 @@ +"""Read-only, bounded analysis of a LoRA's effect on a diffusion model.""" + +from contextlib import closing +import csv +import gc +import heapq +from io import StringIO +import logging + +import folder_paths +import torch +from comfy.weight_adapter import LoRAAdapter +from comfy_api.latest import io +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned + +from .device_utils import cleanup_after_operation +from .lora_alpha import normalize_lora_pair +from .lora_resize import ( + AdapterErrorCapture, build_lora_layer_map, canonical_lora_block_name, + layer_tensor_keys, parse_lora_layers, validate_canonical_blocks, +) +from .merger import _compile_patterns, _matches_any_pattern, load_documentation_from_file +from .model_analysis import ( + CWBStats, RawStats, TensorRecord, _block_name, _common_inputs, + _compute_pair_on_device, _fmt, _is_cuda_oom, _is_float_dtype, _layer_name, + _metrics, _release_cuda_oom, _render_comparison, _render_cwb, +) +from .quantization_guard import inspect_low_bit_input +from .uel_io import stream_work_units + + +def _plan_application(base, lora, base_path, params): + pairs, passthrough = parse_lora_layers(lora.keys()) + validate_canonical_blocks(pairs, "LoRA-on-model analysis") + layer_map = build_lora_layer_map([{"pairs": pairs}], base_path) + targets = {canonical_lora_block_name(key): key for key in base.keys()} + patches = {} + mapped = set() + unmapped_targets = [] + for core, sources in layer_map.items(): + for _, block in sources: + roles = pairs[block] + weight_role = next((role for role in ("set_weight", "diff", "w_norm") if role in roles), None) + specs = [] + if weight_role is not None: + target = targets.get(core) if weight_role == "diff" else None + specs.append((target or targets.get(f"{core}.weight"), weight_role)) + elif "down" in roles and "up" in roles: + specs.append((targets.get(f"{core}.weight"), "pair")) + bias_role = next((role for role in ("diff_b", "b_norm") if role in roles), None) + if bias_role: + specs.append((targets.get(f"{core}.bias"), bias_role)) + for target, role in specs: + if target is not None: + patches.setdefault(target, []).append((block, roles, role)) + mapped.add(block) + else: + unmapped_targets.append(f"{block} ({role}: no base target)") + inventory = {name: [] for name in ("a_only", "b_only", "shape", "nonfloat", "low_bit", "excluded_a", "excluded_b")} + unmatched = sorted(set(pairs) - mapped) + unmapped_targets + patterns = _compile_patterns(params.get("exclude_patterns", ""), glob_mode=params.get("glob_patterns", False)) + low_base = inspect_low_bit_input(base, "Base diffusion model", "LoRA-on-model analysis") + low_lora = inspect_low_bit_input(lora, "LoRA", "LoRA-on-model analysis") + units = [] + for key in sorted(base.keys()): + if _matches_any_pattern(key, patterns, glob_mode=params.get("glob_patterns", False)) != params.get("include_mode", False): + inventory["excluded_a"].append(key) + inventory["excluded_b"].append(key) + continue + specs = patches.get(key, []) + source_keys = list(dict.fromkeys( + source_key for _, roles, role in specs + for source_key in (layer_tensor_keys(roles).values() if role == "pair" else (roles[role],)) + )) + alpha_keys = {roles["alpha"] for _, roles, role in specs if role == "pair" and "alpha" in roles} + if key in low_base or any(source in low_lora and source not in alpha_keys for source in source_keys): + inventory["low_bit"].append(key) + continue + if not _is_float_dtype(base.get_dtype(key)): + inventory["nonfloat"].append(key) + continue + entries = {0: [key]} + if source_keys: + entries[1] = source_keys + units.append((key, entries)) + return units, patches, inventory, unmatched, passthrough + + +def _apply_pair(key, weight, block, roles, loaded, strength): + tensors = {source_key: loaded[(1, source_key)] for source_key in layer_tensor_keys(roles).values()} + down = up = adapter = None + try: + down, up = normalize_lora_pair(tensors[roles["down"]], tensors[roles["up"]], tensors.get(roles.get("alpha")), layer=block) + tensors[roles["down"]], tensors[roles["up"]] = down, up + tensors.pop(roles.get("alpha"), None) + adapter = LoRAAdapter.load(block, tensors, None, tensors.get(roles.get("dora_scale")), set()) + if adapter is None: + raise ValueError(f"LoRA-on-model analysis: could not load adapter {block!r} for {key!r}.") + capture = AdapterErrorCapture(key) + logger = logging.getLogger() + logger.addHandler(capture) + try: + weight = adapter.calculate_weight(weight, key, strength, 1.0, None, lambda value: value, torch.float32, None) + finally: + logger.removeHandler(capture) + if capture.messages: + raise RuntimeError("; ".join(capture.messages)) + return weight + finally: + tensors.clear() + del adapter, down, up, weight + + +def _apply_and_analyze(key, loaded, specs, strength, device, top_count): + source = loaded[(0, key)] + original = patched = delta = None + try: + original = ( + transfer_to_gpu_pinned(source, device=device, dtype=torch.float32) + if str(device).startswith("cuda") else source.to(device=device, dtype=torch.float32) + ) + patched = original.clone() + for block, roles, role in specs: + if role == "pair": + patched = _apply_pair(key, patched, block, roles, loaded, strength) + continue + delta = loaded[(1, roles[role])].to(device=device, dtype=torch.float32) + if delta.shape != patched.shape: + raise ValueError(f"LoRA-on-model analysis: {roles[role]} shape {tuple(delta.shape)} does not match {key} {tuple(patched.shape)}.") + if role == "set_weight": + patched.copy_(delta) + else: + patched.add_(delta, alpha=strength) + delta = None + if patched.shape != original.shape: + raise ValueError(f"LoRA-on-model analysis: applying {key!r} changed its shape; it cannot be compared with the base tensor.") + return _compute_pair_on_device(original, patched, device, False, top_count) + finally: + del source, original, patched, delta + + +def _analyze_unit(key, loaded, specs, strength, device, top_count): + try: + return (*_apply_and_analyze(key, loaded, specs, strength, device, top_count), False) + except Exception as error: + if not str(device).startswith("cuda") or not _is_cuda_oom(error): + raise + error.__traceback__ = None + _release_cuda_oom() + try: + return (*_apply_and_analyze(key, loaded, specs, strength, "cpu", top_count), True) + except Exception as cpu_error: + raise RuntimeError(f"LoRA-on-model analysis CPU fallback failed for {key!r}: {cpu_error}") from cpu_error + + +def _csv(headers, rows): + output = StringIO(newline="") + writer = csv.writer(output) + writer.writerow(headers) + writer.writerows(rows) + return output.getvalue() + + +def analyze_lora_on_model(model_a, model_b, params): + lora_path = folder_paths.get_full_path_or_raise("loras", model_a) + base_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + strength = params.get("strength", 1.0) + records, fallback_keys, top_heap = [], [], [] + global_raw = RawStats() + layer_stats, block_stats, block_cwb = {}, {}, {} + try: + with MemoryEfficientSafeOpen(base_path, low_memory=True) as base, MemoryEfficientSafeOpen(lora_path, low_memory=True) as lora: + units, patches, inventory, unmatched, passthrough = _plan_application(base, lora, base_path, params) + device = params["process_device"] + top_limit = params["top_weight_differences"] + with closing(stream_work_units({0: base, 1: lora}, units, pin_memory=str(device).startswith("cuda"))) as stream: + for key, loaded in stream: + try: + raw, cwb, top, fallback = _analyze_unit(key, loaded, patches.get(key, []), strength, device, top_limit) + shape = tuple(base.get_shape(key)) + records.append(TensorRecord(key, shape, str(base.get_dtype(key)), str(torch.float32), raw, cwb, fallback)) + global_raw.add(raw) + layer_stats.setdefault(_layer_name(key), RawStats()).add(raw) + block_stats.setdefault(_block_name(key), RawStats()).add(raw) + block_cwb.setdefault(_block_name(key), CWBStats()).add(cwb) + if fallback: + fallback_keys.append(key) + for value, index, a, b in top: + entry = (value, key, index, a, b, shape) + if len(top_heap) < top_limit: + heapq.heappush(top_heap, entry) + elif top_limit and value > top_heap[0][0]: + heapq.heapreplace(top_heap, entry) + finally: + loaded.clear() + if params["force_clear_cache"]: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + label = f"{model_b} + {model_a} (strength={strength})" + comparison = _render_comparison(model_b, label, "LoRA applied to diffusion model", records, global_raw, layer_stats, block_stats, inventory, top_heap, fallback_keys, base.metadata(), base.metadata()) + comparison = ( + "# LoRA-on-model analysis\n\nModel A input is the LoRA; Model B input is the base diffusion model. " + "The comparison below uses original base as A and patched base as B; delta = patched − original. " + "Unmodified base tensors are included unless filtered. No model is saved.\n\n" + + comparison + + "\n## Unmapped LoRA groups\n\n" + ("\n".join(f"- `{name}`" for name in unmatched) or "None.") + + "\n\n## Unused LoRA tensors\n\n" + ("\n".join(f"- `{name}`" for name in passthrough) or "None.") + "\n" + ) + cwb_report = _render_cwb(model_b, label, records, block_cwb, fallback_keys, "original versus LoRA-patched base tensor rows (index only)") + metric_names = ("mae", "mse", "rmse", "max", "relative_l2", "cosine", "pearson", "exact", "sign", "norm_ratio", "coverage") + metrics_csv = _csv(["key", "parameters", *metric_names], ([record.key, record.raw.total, *(_fmt(_metrics(record.raw)[name]) for name in metric_names)] for record in records)) + cwb_names = ("pair_mean", "pair_min", "pair_max") + cwb_csv = _csv(["key", *cwb_names], ([record.key, *(_fmt(record.cwb.values()[name]) for name in cwb_names)] for record in records)) + return comparison, cwb_report, metrics_csv, cwb_csv + finally: + cleanup_after_operation() + + +class LoRAOnModelAnalysis(io.ComfyNode): + @classmethod + def define_schema(cls): + inputs = _common_inputs("diffusion_models", alignment_control=False) + inputs[1] = io.Combo.Input("model_a", options=folder_paths.get_filename_list("loras"), tooltip="LoRA to apply transiently to Model B before comparing the original and patched base tensors.") + inputs[2] = io.Combo.Input("model_b", options=folder_paths.get_filename_list("diffusion_models"), tooltip="Base diffusion model. Its original weights are compared against a transient LoRA-patched copy, one tensor at a time.") + inputs.append(io.Float.Input("strength", default=1.0, min=-10.0, max=10.0, step=0.01, tooltip="LoRA application strength. Alpha scaling and DoRA reconstruction use ComfyUI's adapter; explicit set_weight patches retain their existing replacement semantics.")) + return io.Schema( + node_id="LoRAOnModelAnalysis", display_name="Analyze LoRA Effect on Diffusion Model", + category="ModelUtils/Analysis", + description="Compare the base diffusion model with its LoRA-patched weights using bounded streaming. No model is written or inferred.", + inputs=inputs, + outputs=[io.String.Output(display_name="comparison_report"), io.String.Output(display_name="cwb_report"), io.String.Output(display_name="documentation"), io.String.Output(display_name="layerwise_metrics_csv"), io.String.Output(display_name="layerwise_cwb_csv")], + ) + + @classmethod + def execute(cls, **kwargs): + docs = load_documentation_from_file("lora_model_analysis.md") + if kwargs["execution_mode"] == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No comparison performed.", "No CWB analysis performed.", docs, "", "") + comparison, cwb, metrics_csv, cwb_csv = analyze_lora_on_model(kwargs["model_a"], kwargs["model_b"], kwargs) + return io.NodeOutput(comparison, cwb, docs, metrics_csv, cwb_csv) diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index cafe021..533f639 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -38,6 +38,20 @@ from .lora_extract_svd import ( MIN_SV = 1e-6 + +class AdapterErrorCapture(logging.Handler): + """Observe errors ComfyUI logs and swallows for one adapter target.""" + + def __init__(self, target_key): + super().__init__(level=logging.ERROR) + self.target_key = target_key + self.messages = [] + + def emit(self, record): + message = record.getMessage() + if message.startswith("ERROR ") and self.target_key in message: + self.messages.append(message) + LORA_PAIR_SUFFIXES = ( (".lora_A.default.weight", ".lora_B.default.weight", ".alpha", "peft"), (".lora_down.weight", ".lora_up.weight", ".alpha", "comfy"), @@ -1698,19 +1712,6 @@ def merge_loras_to_model( ) return f"L{info['index'] + 1}::{block_name}\n{sources}\n{details}" - class AdapterErrorCapture(logging.Handler): - """Observe errors ComfyUI logs and swallows for one adapter target.""" - - def __init__(self, target_key): - super().__init__(level=logging.ERROR) - self.target_key = target_key - self.messages = [] - - def emit(self, record): - message = record.getMessage() - if message.startswith("ERROR ") and self.target_key in message: - self.messages.append(message) - stats = { "merged": 0, "base_bias_retained": 0, diff --git a/tests/test_groups.toml b/tests/test_groups.toml index 316d0f1..f533741 100644 --- a/tests/test_groups.toml +++ b/tests/test_groups.toml @@ -26,6 +26,10 @@ python_tests = [ paths = ["nodes/model_analysis.py", "run_model_analysis_api_workflow.py"] python_tests = ["tests/test_model_analysis.py", "tests/test_model_analysis_api_runner.py"] +[groups.lora_model_analysis] +paths = ["nodes/lora_model_analysis.py", "nodes/lora_resize.py", "nodes/model_analysis.py"] +python_tests = ["tests/test_lora_model_analysis.py"] + [groups.adaptive_extraction] paths = [ "nodes/adaptive_svd.py", "nodes/dora_extract_wd.py", "nodes/dora_learned_wd.py", diff --git a/tests/test_lora_model_analysis.py b/tests/test_lora_model_analysis.py new file mode 100644 index 0000000..d9fd754 --- /dev/null +++ b/tests/test_lora_model_analysis.py @@ -0,0 +1,234 @@ +import csv +import importlib +from io import StringIO +import logging +from pathlib import Path +import sys +import types + +import pytest +import torch +from unifiedefficientloader import IncrementalSafetensorsWriter + + +@pytest.fixture(scope="module") +def analysis(): + name = "modelutils_lora_on_model_tests" + package = types.ModuleType(name) + package.__path__ = [str(Path(__file__).resolve().parents[1])] + sys.modules[name] = package + return importlib.import_module(f"{name}.nodes.lora_model_analysis") + + +def write(path, tensors): + with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer: + writer.write_batch(list(tensors.items())) + + +def setup(monkeypatch, tmp_path, analysis, adapter, base=None): + paths = {"loras": tmp_path / "adapter.safetensors", "diffusion_models": tmp_path / "base.safetensors"} + write(paths["loras"], adapter) + write(paths["diffusion_models"], base if base is not None else { + "model.diffusion_model.layer.weight": torch.ones((2, 2)), + "model.diffusion_model.untouched.weight": torch.ones((2, 2)), + }) + monkeypatch.setattr(analysis.folder_paths, "get_full_path_or_raise", lambda category, name: str(paths[category])) + monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None) + return paths + + +def params(**overrides): + return {"execution_mode": "ANALYZE", "process_device": "cpu", "top_weight_differences": 3, "force_clear_cache": False, "exclude_patterns": "", "glob_patterns": False, "include_mode": False, "strength": 1.0, **overrides} + + +def pair(alpha=True): + result = { + "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]), + "diffusion_model.layer.lora_B.weight": torch.tensor([[2.0], [4.0]]), + } + if alpha: + result["diffusion_model.layer.alpha"] = torch.tensor(0.5) + return result + + +@pytest.mark.parametrize("alpha,strength,mae,maximum", [(True, 1.0, 2.25, 4), (False, 1.0, 4.5, 8), (True, 0.5, 1.125, 2), (True, 0.0, 0, 0)]) +def test_reports_actual_applied_delta_without_saving(monkeypatch, tmp_path, analysis, alpha, strength, mae, maximum): + paths = setup(monkeypatch, tmp_path, analysis, pair(alpha)) + before = {path: path.read_bytes() for path in paths.values()} + comparison, cwb, metrics_csv, _ = analysis.analyze_lora_on_model("adapter", "base", params(strength=strength)) + rows = {row["key"]: row for row in csv.DictReader(StringIO(metrics_csv))} + row = rows["model.diffusion_model.layer.weight"] + assert float(row["mae"]) == pytest.approx(mae) + assert float(row["max"]) == maximum + assert float(rows["model.diffusion_model.untouched.weight"]["mae"]) == 0 + assert "patched − original" in comparison + assert "original versus LoRA-patched" in cwb + assert {path: path.read_bytes() for path in paths.values()} == before + assert set(tmp_path.iterdir()) == set(paths.values()) + + +@pytest.mark.parametrize("suffixes", [(".lora_down.weight", ".lora_up.weight"), (".lora_A.default.weight", ".lora_B.default.weight"), (".lora_A", ".lora_B")]) +def test_pair_formats_and_flattened_reference_mapping(monkeypatch, tmp_path, analysis, suffixes): + setup(monkeypatch, tmp_path, analysis, { + "lora_unet_layer" + suffixes[0]: torch.ones((1, 2)), + "lora_unet_layer" + suffixes[1]: torch.ones((2, 1)), + }) + _, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params()) + rows = list(csv.DictReader(StringIO(text))) + assert float(rows[0]["mae"]) == 1 + + +def test_direct_weight_bias_replacement_and_unmapped_are_explicit(monkeypatch, tmp_path, analysis): + setup(monkeypatch, tmp_path, analysis, { + "diffusion_model.layer.diff": torch.full((2, 2), 2.0), + "diffusion_model.layer.diff_b": torch.full((2,), 3.0), + "diffusion_model.replaced.set_weight": torch.full((2, 2), 5.0), + "diffusion_model.missing.diff": torch.ones((2, 2)), + }, { + "diffusion_model.layer.weight": torch.ones((2, 2)), + "diffusion_model.layer.bias": torch.ones(2), + "diffusion_model.replaced.weight": torch.ones((2, 2)), + }) + report, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params(strength=0.5)) + rows = {row["key"]: row for row in csv.DictReader(StringIO(text))} + assert float(rows["diffusion_model.layer.weight"]["mae"]) == 1 + assert float(rows["diffusion_model.layer.bias"]["mae"]) == 1.5 + assert float(rows["diffusion_model.replaced.weight"]["mae"]) == 4 + assert "diffusion_model.missing" in report + + +def test_bounded_stream_release_and_filtering(monkeypatch, tmp_path, analysis): + setup(monkeypatch, tmp_path, analysis, pair()) + calls, released = [], [] + original_stream = analysis.MemoryEfficientSafeOpen.async_stream + original_mark = analysis.MemoryEfficientSafeOpen.mark_processed + + def stream(handler, keys, **kwargs): + calls.append((list(keys), kwargs)) + return original_stream(handler, keys, **kwargs) + + def mark(handler, key): + released.append(key) + return original_mark(handler, key) + + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream) + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", mark) + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "get_tensor", lambda *args: pytest.fail("synchronous tensor load")) + _, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params(exclude_patterns="*layer.weight", glob_patterns=True, include_mode=True)) + assert len(list(csv.DictReader(StringIO(text)))) == 1 + assert len(calls) == 2 + assert all(kwargs == {"batch_size": 1, "prefetch_batches": 1, "pin_memory": False} for _, kwargs in calls) + assert sorted(released) == sorted(key for keys, _ in calls for key in keys) + assert len(released) == 4 + + +def test_logged_adapter_failure_is_not_reported_as_zero_delta(monkeypatch, tmp_path, analysis): + setup(monkeypatch, tmp_path, analysis, pair()) + marks = [] + original_mark = analysis.MemoryEfficientSafeOpen.mark_processed + + def mark(handler, key): + marks.append(key) + return original_mark(handler, key) + + def fail(self, weight, key, *args): + logging.error("ERROR LoRA %s deliberately failed", key) + return weight + + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", mark) + monkeypatch.setattr(analysis.LoRAAdapter, "calculate_weight", fail) + with pytest.raises(RuntimeError, match="deliberately failed"): + analysis.analyze_lora_on_model("adapter", "base", params()) + assert len(marks) == 4 + assert not any(isinstance(handler, analysis.AdapterErrorCapture) for handler in logging.getLogger().handlers) + + +def test_complete_unit_cuda_retry(monkeypatch, analysis): + devices = [] + + def calculate(key, loaded, specs, strength, device, top_count): + devices.append(device) + if device == "cuda": + raise torch.cuda.OutOfMemoryError("test") + return analysis.RawStats(), analysis.CWBStats(), [] + + monkeypatch.setattr(analysis, "_apply_and_analyze", calculate) + monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None) + assert analysis._analyze_unit("layer", {}, [], 1, "cuda", 0)[-1] is True + assert devices == ["cuda", "cpu"] + + +def test_dora_uses_reconstructed_weight_not_raw_factor_delta(monkeypatch, tmp_path, analysis): + setup(monkeypatch, tmp_path, analysis, { + "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 0.0]]), + "diffusion_model.layer.lora_B.weight": torch.tensor([[1.0], [0.0]]), + "diffusion_model.layer.alpha": torch.tensor(0.5), + "diffusion_model.layer.dora_scale": torch.tensor([3.0, 4.0]), + }, {"diffusion_model.layer.weight": torch.eye(2)}) + _, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params()) + row = next(csv.DictReader(StringIO(text))) + # Installed Comfy output-axis DoRA scales by the original base row norms. + # The resulting diagonal is (4.5, 4), not the raw-factor result (1.5, 1). + assert float(row["mae"]) == pytest.approx(1.625) + assert float(row["max"]) == pytest.approx(3.5) + + +def test_locon_mid_and_scalar_low_bit_alpha(monkeypatch, tmp_path, analysis): + setup(monkeypatch, tmp_path, analysis, { + "diffusion_model.layer.lora_down.weight": torch.ones((1, 1, 1, 1)), + "diffusion_model.layer.lora_up.weight": torch.ones((1, 1, 1, 1)), + "diffusion_model.layer.lora_mid.weight": torch.ones((1, 1, 3, 3)), + "diffusion_model.layer.alpha": torch.tensor(1, dtype=torch.uint8), + }, {"diffusion_model.layer.weight": torch.zeros((1, 1, 3, 3))}) + _, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params()) + row = next(csv.DictReader(StringIO(text))) + assert float(row["mae"]) == 1 + + +def test_low_bit_factors_reported_without_materializing(monkeypatch, tmp_path, analysis): + tensors = pair() + tensors["diffusion_model.layer.lora_A.weight"] = torch.ones((1, 2), dtype=torch.uint8) + setup(monkeypatch, tmp_path, analysis, tensors) + requested = [] + original = analysis.MemoryEfficientSafeOpen.async_stream + + def stream(handler, keys, **kwargs): + requested.extend(keys) + return original(handler, keys, **kwargs) + + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream) + report, _, text, _ = analysis.analyze_lora_on_model("adapter", "base", params()) + assert requested == ["model.diffusion_model.untouched.weight"] + assert "model.diffusion_model.layer.weight" in report.replace(r"\_", "_") + assert [row["key"] for row in csv.DictReader(StringIO(text))] == ["model.diffusion_model.untouched.weight"] + + +def test_streams_close_after_failure(monkeypatch, tmp_path, analysis): + setup(monkeypatch, tmp_path, analysis, pair()) + streams_closed = [] + original = analysis.MemoryEfficientSafeOpen.async_stream + + def stream(handler, keys, **kwargs): + iterator = original(handler, keys, **kwargs) + try: + yield from iterator + finally: + iterator.close() + streams_closed.append(True) + + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream) + monkeypatch.setattr(analysis.LoRAAdapter, "calculate_weight", lambda *args: (_ for _ in ()).throw(ValueError("test failure"))) + with pytest.raises(ValueError, match="test failure"): + analysis.analyze_lora_on_model("adapter", "base", params()) + assert len(streams_closed) == 2 + + +def test_schema_and_documentation_mode(monkeypatch, analysis): + monkeypatch.setattr(analysis.folder_paths, "get_filename_list", lambda category: [category]) + schema = analysis.LoRAOnModelAnalysis.define_schema() + assert schema.inputs[1].id == "model_a" and schema.inputs[1].options == ["loras"] + assert schema.inputs[2].id == "model_b" and schema.inputs[2].options == ["diffusion_models"] + assert len(schema.outputs) == 5 + assert schema.inputs[-1].id == "strength" and schema.inputs[-1].default == 1 + monkeypatch.setattr(analysis, "analyze_lora_on_model", lambda *args: pytest.fail("documentation loaded tensors")) + assert analysis.LoRAOnModelAnalysis.execute(execution_mode="DOCUMENTATION ONLY")[2].startswith("# LoRA-on-model analysis")