235 lines
11 KiB
Python
235 lines
11 KiB
Python
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")
|