EthanfelandClaude Opus 4.6
2bc4af3580
Single-LoRA prefixes produce identical patches regardless of merge config (forced to weighted_sum, early return at _merge_diffs). Cache results keyed by (label_prefix, auto_strength) to skip diff recomputation for candidates 2+ sharing the same auto_strength value. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2443 lines
110 KiB
Python
2443 lines
110 KiB
Python
import importlib.util
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
try:
|
|
import torch
|
|
except ModuleNotFoundError:
|
|
torch = None
|
|
|
|
|
|
def _install_stubs():
|
|
tmpdir = tempfile.gettempdir()
|
|
|
|
folder_paths = types.ModuleType("folder_paths")
|
|
folder_paths.models_dir = tmpdir
|
|
folder_paths.add_model_folder_path = lambda *args, **kwargs: None
|
|
folder_paths.get_temp_directory = lambda: tmpdir
|
|
folder_paths.get_user_directory = lambda: tmpdir
|
|
folder_paths.get_folder_paths = lambda _kind: [tmpdir]
|
|
folder_paths.get_filename_list = lambda _kind: []
|
|
folder_paths.get_full_path_or_raise = lambda _kind, name: name
|
|
folder_paths.get_full_path = lambda _kind, name: name
|
|
sys.modules["folder_paths"] = folder_paths
|
|
|
|
comfy = types.ModuleType("comfy")
|
|
utils = types.ModuleType("comfy.utils")
|
|
|
|
def get_attr(obj, path):
|
|
for part in path.split("."):
|
|
obj = getattr(obj, part)
|
|
return obj
|
|
|
|
class ProgressBar:
|
|
def __init__(self, total):
|
|
self.total = total
|
|
self.value = 0
|
|
|
|
def update(self, amount):
|
|
self.value += amount
|
|
|
|
utils.get_attr = get_attr
|
|
utils.load_torch_file = lambda _path, safe_load=True: {}
|
|
utils.ProgressBar = ProgressBar
|
|
|
|
sd = types.ModuleType("comfy.sd")
|
|
sd.load_lora_for_models = lambda model, clip, lora_dict, model_strength, clip_strength: (model, clip)
|
|
|
|
lora = types.ModuleType("comfy.lora")
|
|
lora.model_lora_keys_unet = lambda model, mapping: {}
|
|
lora.model_lora_keys_clip = lambda clip, mapping: {}
|
|
|
|
model_management = types.ModuleType("comfy.model_management")
|
|
model_management.get_free_memory = lambda _device: 0
|
|
|
|
weight_adapter = types.ModuleType("comfy.weight_adapter")
|
|
weight_adapter_lora = types.ModuleType("comfy.weight_adapter.lora")
|
|
weight_adapter_lokr = types.ModuleType("comfy.weight_adapter.lokr")
|
|
weight_adapter_loha = types.ModuleType("comfy.weight_adapter.loha")
|
|
|
|
class LoRAAdapter:
|
|
def __init__(self, loaded_keys, weights):
|
|
self.loaded_keys = loaded_keys
|
|
self.weights = weights
|
|
|
|
class LoKrAdapter:
|
|
def __init__(self, loaded_keys, weights):
|
|
self.loaded_keys = loaded_keys
|
|
self.weights = weights
|
|
|
|
class LoHaAdapter:
|
|
def __init__(self, loaded_keys, weights):
|
|
self.loaded_keys = loaded_keys
|
|
self.weights = weights
|
|
|
|
weight_adapter_lora.LoRAAdapter = LoRAAdapter
|
|
weight_adapter_lokr.LoKrAdapter = LoKrAdapter
|
|
weight_adapter_loha.LoHaAdapter = LoHaAdapter
|
|
|
|
comfy.utils = utils
|
|
comfy.sd = sd
|
|
comfy.lora = lora
|
|
comfy.model_management = model_management
|
|
|
|
sys.modules["comfy"] = comfy
|
|
sys.modules["comfy.utils"] = utils
|
|
sys.modules["comfy.sd"] = sd
|
|
sys.modules["comfy.lora"] = lora
|
|
sys.modules["comfy.model_management"] = model_management
|
|
sys.modules["comfy.weight_adapter"] = weight_adapter
|
|
sys.modules["comfy.weight_adapter.lora"] = weight_adapter_lora
|
|
sys.modules["comfy.weight_adapter.lokr"] = weight_adapter_lokr
|
|
sys.modules["comfy.weight_adapter.loha"] = weight_adapter_loha
|
|
|
|
safetensors = types.ModuleType("safetensors")
|
|
safetensors.safe_open = mock.MagicMock()
|
|
safetensors_torch = types.ModuleType("safetensors.torch")
|
|
safetensors_torch.save_file = lambda state_dict, path, metadata=None: None
|
|
safetensors.torch = safetensors_torch
|
|
sys.modules["safetensors"] = safetensors
|
|
sys.modules["safetensors.torch"] = safetensors_torch
|
|
|
|
|
|
def _load_module():
|
|
"""Load lora_optimizer module with stubs in place."""
|
|
_install_stubs()
|
|
spec = importlib.util.spec_from_file_location(
|
|
"lora_optimizer",
|
|
os.path.join(os.path.dirname(__file__), "..", "lora_optimizer.py"),
|
|
)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
return mod
|
|
|
|
|
|
if torch is not None:
|
|
_install_stubs()
|
|
MODULE_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "lora_optimizer.py")
|
|
SPEC = importlib.util.spec_from_file_location("lora_optimizer_under_test", MODULE_PATH)
|
|
lora_optimizer = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(lora_optimizer)
|
|
else:
|
|
lora_optimizer = None
|
|
|
|
|
|
def _make_model():
|
|
layer = types.SimpleNamespace(weight=torch.zeros(1, 1))
|
|
return types.SimpleNamespace(model=types.SimpleNamespace(layer=layer))
|
|
|
|
|
|
def _make_lora_entry(prefix_to_value, strength=1.0, clip_strength=None, key_filter="all", conflict_mode="all", name="demo"):
|
|
lora = {}
|
|
for prefix, value in prefix_to_value.items():
|
|
lora[f"{prefix}.lora_up.weight"] = torch.tensor([[float(value)]], dtype=torch.float32)
|
|
lora[f"{prefix}.lora_down.weight"] = torch.tensor([[1.0]], dtype=torch.float32)
|
|
return {
|
|
"name": name,
|
|
"lora": lora,
|
|
"strength": strength,
|
|
"clip_strength": clip_strength,
|
|
"key_filter": key_filter,
|
|
"conflict_mode": conflict_mode,
|
|
}
|
|
|
|
|
|
@unittest.skipIf(torch is None, "torch is not installed in this environment")
|
|
class LoRAOptimizerTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.optimizer = lora_optimizer.LoRAOptimizer()
|
|
self.model = _make_model()
|
|
|
|
def test_lora_format_cache_avoids_repeated_detection(self):
|
|
"""After detecting a LoRA's format once, subsequent prefixes should reuse it."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
lora_dict = {
|
|
"unet.a.lora_B.weight": torch.tensor([[1.0]], dtype=torch.float32),
|
|
"unet.a.lora_A.weight": torch.tensor([[1.0]], dtype=torch.float32),
|
|
"unet.b.lora_B.weight": torch.tensor([[2.0]], dtype=torch.float32),
|
|
"unet.b.lora_A.weight": torch.tensor([[1.0]], dtype=torch.float32),
|
|
}
|
|
result1 = optimizer._get_lora_key_info(lora_dict, "unet.a")
|
|
self.assertIsNotNone(result1)
|
|
self.assertIn(id(lora_dict), optimizer._lora_format_cache)
|
|
result2 = optimizer._get_lora_key_info(lora_dict, "unet.b")
|
|
self.assertIsNotNone(result2)
|
|
|
|
def test_target_groups_merge_aliases_for_same_target(self):
|
|
groups = self.optimizer._build_target_groups(
|
|
["alias_a", "alias_b", "other"],
|
|
{"alias_a": "layer.weight", "alias_b": "layer.weight", "other": "other.weight"},
|
|
{},
|
|
)
|
|
|
|
self.assertEqual(set(groups.keys()), {"alias_a", "other"})
|
|
self.assertEqual(groups["alias_a"]["aliases"], ["alias_a", "alias_b"])
|
|
|
|
def test_group_analysis_detects_alias_overlap(self):
|
|
active_loras = [
|
|
_make_lora_entry({"alias_a": 1.0}, name="A"),
|
|
_make_lora_entry({"alias_b": -1.0}, name="B"),
|
|
]
|
|
target_groups = self.optimizer._build_target_groups(
|
|
["alias_a", "alias_b"],
|
|
{"alias_a": "layer.weight", "alias_b": "layer.weight"},
|
|
{},
|
|
)
|
|
|
|
analysis = self.optimizer._run_group_analysis(
|
|
target_groups, active_loras, self.model, None, torch.device("cpu")
|
|
)
|
|
|
|
self.assertEqual(analysis["prefix_count"], 1)
|
|
stats = analysis["prefix_stats"]["alias_a"]
|
|
self.assertEqual(stats["n_loras"], 2)
|
|
self.assertGreater(stats["conflict_ratio"], 0.99)
|
|
|
|
def test_same_lora_aliases_are_aggregated_before_analysis(self):
|
|
target_group = {
|
|
"target_key": "layer.weight",
|
|
"is_clip": False,
|
|
"aliases": ["alias_a", "alias_b"],
|
|
"label_prefix": "alias_a",
|
|
}
|
|
active_loras = [
|
|
_make_lora_entry({"alias_a": 1.0, "alias_b": 2.0}, name="A"),
|
|
]
|
|
|
|
prepared = self.optimizer._prepare_group_diffs(
|
|
target_group, active_loras, self.model, None, torch.device("cpu")
|
|
)
|
|
|
|
self.assertAlmostEqual(prepared["diffs"][0].item(), 3.0)
|
|
|
|
def test_exact_linear_patch_matches_dense_sum(self):
|
|
target_group = {
|
|
"target_key": "layer.weight",
|
|
"is_clip": False,
|
|
"aliases": ["alias_a", "alias_b"],
|
|
"label_prefix": "alias_a",
|
|
}
|
|
active_loras = [
|
|
_make_lora_entry({"alias_a": 1.0}, name="A"),
|
|
_make_lora_entry({"alias_b": 2.0}, name="B"),
|
|
]
|
|
|
|
patch_info = self.optimizer._build_exact_linear_patch(
|
|
target_group, active_loras, raw_n_loras=2, mode="weighted_sum"
|
|
)
|
|
|
|
diff = self.optimizer._expand_patch_to_diff(patch_info["patch"])
|
|
self.assertAlmostEqual(diff.item(), 3.0)
|
|
|
|
def test_expand_patch_to_diff_supports_lokr_and_loha(self):
|
|
lokr_patch = lora_optimizer.LoKrAdapter(
|
|
set(),
|
|
(
|
|
torch.tensor([[2.0]], dtype=torch.float32),
|
|
torch.tensor([[3.0]], dtype=torch.float32),
|
|
1.0,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
),
|
|
)
|
|
loha_patch = lora_optimizer.LoHaAdapter(
|
|
set(),
|
|
(
|
|
torch.tensor([[2.0]], dtype=torch.float32),
|
|
torch.tensor([[3.0]], dtype=torch.float32),
|
|
1.0,
|
|
torch.tensor([[4.0]], dtype=torch.float32),
|
|
torch.tensor([[5.0]], dtype=torch.float32),
|
|
None,
|
|
None,
|
|
None,
|
|
),
|
|
)
|
|
|
|
self.assertAlmostEqual(self.optimizer._expand_patch_to_diff(lokr_patch).item(), 6.0)
|
|
self.assertAlmostEqual(self.optimizer._expand_patch_to_diff(loha_patch).item(), 120.0)
|
|
|
|
def test_auto_strength_uses_exact_streamed_energy(self):
|
|
active_loras = [
|
|
{"name": "A", "strength": 1.0, "clip_strength": None},
|
|
{"name": "B", "strength": 1.0, "clip_strength": None},
|
|
]
|
|
branch_energy = {
|
|
"model": {
|
|
"norm_sq": [1.0, 1.0],
|
|
"dot": {(0, 1): 1.0},
|
|
},
|
|
"clip": {
|
|
"norm_sq": [0.0, 0.0],
|
|
"dot": {(0, 1): 0.0},
|
|
},
|
|
}
|
|
|
|
info = self.optimizer._compute_auto_strengths(active_loras, branch_energy)
|
|
self.assertAlmostEqual(info["model_scale"], 0.5)
|
|
self.assertAlmostEqual(info["model_strengths"][0], 0.5)
|
|
self.assertAlmostEqual(info["model_strengths"][1], 0.5)
|
|
|
|
def test_pair_metrics_capture_excess_conflict_and_subspace_overlap(self):
|
|
diff_a = torch.tensor([[1.0, 0.0], [0.0, 0.0]], dtype=torch.float32)
|
|
diff_b = torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
|
|
basis_a = self.optimizer._compute_subspace_basis(diff_a, rank_hint=1)
|
|
basis_b = self.optimizer._compute_subspace_basis(diff_b, rank_hint=1)
|
|
|
|
metrics = self.optimizer._sample_pair_metrics(diff_a, diff_b, basis_a=basis_a, basis_b=basis_b)
|
|
|
|
self.assertEqual(metrics["overlap"], 0)
|
|
self.assertAlmostEqual(metrics["subspace_overlap"], 0.0, places=4)
|
|
self.assertAlmostEqual(metrics["excess_conflict"], 0.0, places=4)
|
|
|
|
def test_block_smoothing_populates_decision_metrics(self):
|
|
prefix_stats = {
|
|
"block_0.attn.q": {
|
|
"n_loras": 2,
|
|
"conflict_ratio": 0.10,
|
|
"excess_conflict": 0.10,
|
|
"avg_cos_sim": 0.20,
|
|
"avg_subspace_overlap": 0.30,
|
|
"magnitude_ratio": 1.0,
|
|
"per_lora_norm_sq": {0: 1.0, 1: 1.0},
|
|
},
|
|
"block_0.attn.k": {
|
|
"n_loras": 2,
|
|
"conflict_ratio": 0.50,
|
|
"excess_conflict": 0.50,
|
|
"avg_cos_sim": 0.20,
|
|
"avg_subspace_overlap": 0.30,
|
|
"magnitude_ratio": 1.0,
|
|
"per_lora_norm_sq": {0: 1.0, 1: 1.0},
|
|
},
|
|
}
|
|
|
|
smoothed = self.optimizer._apply_block_smoothing(prefix_stats, strength=0.5)
|
|
|
|
self.assertIn("decision_conflict", smoothed["block_0.attn.q"])
|
|
self.assertGreater(smoothed["block_0.attn.q"]["decision_conflict"], 0.10)
|
|
self.assertLess(smoothed["block_0.attn.q"]["decision_conflict"], 0.50)
|
|
self.assertEqual(smoothed["block_0.attn.q"]["block_name"], smoothed["block_0.attn.k"]["block_name"])
|
|
|
|
def test_auto_select_uses_excess_conflict_and_subspace(self):
|
|
mode, _density, _sign, _reasoning = self.optimizer._auto_select_params(
|
|
0.55, 1.0, avg_cos_sim=0.0,
|
|
avg_excess_conflict=0.05, avg_subspace_overlap=0.10,
|
|
)
|
|
self.assertEqual(mode, "weighted_average")
|
|
|
|
mode, _density, _sign, _reasoning = self.optimizer._auto_select_params(
|
|
0.55, 1.0, avg_cos_sim=0.15,
|
|
avg_excess_conflict=0.40, avg_subspace_overlap=0.85,
|
|
)
|
|
self.assertEqual(mode, "ties")
|
|
|
|
def test_python_evaluator_spec_and_runner(self):
|
|
builder = lora_optimizer.BuildAutoTunerPythonEvaluator()
|
|
evaluator, = builder.build(
|
|
module_path=os.path.join(tempfile.gettempdir(), "dummy_eval.py"),
|
|
callable_name="evaluate_candidate",
|
|
combine_mode="blend",
|
|
weight=0.7,
|
|
context_json='{"prompt":"test"}',
|
|
)
|
|
|
|
with open(evaluator["module_path"], "w") as f:
|
|
f.write(
|
|
"def evaluate_candidate(model=None, clip=None, lora_data=None, config=None, context=None, analysis_summary=None):\n"
|
|
" assert context['prompt'] == 'test'\n"
|
|
" return {'score': 0.75, 'details': {'ok': True}}\n"
|
|
)
|
|
|
|
result = lora_optimizer._run_autotuner_evaluator(
|
|
evaluator,
|
|
model=self.model,
|
|
clip=None,
|
|
lora_data={"model_patches": {}, "clip_patches": {}},
|
|
config={"merge_mode": "weighted_average"},
|
|
analysis_summary={"avg_conflict_ratio": 0.1},
|
|
)
|
|
|
|
self.assertAlmostEqual(result["score"], 0.75)
|
|
self.assertEqual(result["details"]["ok"], True)
|
|
|
|
def test_conflict_editor_preserves_key_filter_for_tuple_stacks(self):
|
|
editor = lora_optimizer.LoRAConflictEditor()
|
|
editor.loaded_loras["demo"] = {
|
|
"alias_a.lora_up.weight": torch.tensor([[1.0]], dtype=torch.float32),
|
|
"alias_a.lora_down.weight": torch.tensor([[1.0]], dtype=torch.float32),
|
|
}
|
|
|
|
enriched, _report, _strategy = editor.analyze_and_enrich(
|
|
[("demo", 1.0, 1.0, "all", "shared_only")],
|
|
"auto",
|
|
conflict_mode_1="auto",
|
|
)
|
|
|
|
self.assertEqual(len(enriched[0]), 5)
|
|
self.assertEqual(enriched[0][4], "shared_only")
|
|
|
|
def test_save_merged_lora_uses_canonical_prefix(self):
|
|
saver = lora_optimizer.SaveMergedLoRA()
|
|
patch = lora_optimizer.LoRAAdapter(
|
|
set(),
|
|
(torch.tensor([[1.0]]), torch.tensor([[1.0]]), 1.0, None, None, None),
|
|
)
|
|
captured = {}
|
|
|
|
with mock.patch.object(lora_optimizer, "save_file", side_effect=lambda state_dict, path, metadata=None: captured.update({"state_dict": state_dict, "path": path, "metadata": metadata})):
|
|
save_path, = saver.save_lora(
|
|
{
|
|
"model_patches": {"layer.weight": patch},
|
|
"clip_patches": {},
|
|
"key_map": {
|
|
"layer.weight": {
|
|
"canonical_prefix": "canonical_alias",
|
|
"aliases": ["alias_a", "canonical_alias"],
|
|
}
|
|
},
|
|
"output_strength": 1.0,
|
|
"clip_strength": 1.0,
|
|
},
|
|
tempfile.gettempdir(),
|
|
"merged_test",
|
|
save_rank=0,
|
|
bake_strength=False,
|
|
)
|
|
|
|
self.assertTrue(save_path.endswith(".safetensors"))
|
|
self.assertIn("canonical_alias.lora_up.weight", captured["state_dict"])
|
|
|
|
def test_save_nodes_block_directory_traversal(self):
|
|
merged_saver = lora_optimizer.SaveMergedLoRA()
|
|
tuner_saver = lora_optimizer.SaveTunerData()
|
|
patch = lora_optimizer.LoRAAdapter(
|
|
set(),
|
|
(torch.tensor([[1.0]]), torch.tensor([[1.0]]), 1.0, None, None, None),
|
|
)
|
|
|
|
with self.assertRaises(ValueError):
|
|
merged_saver.save_lora(
|
|
{
|
|
"model_patches": {"layer.weight": patch},
|
|
"clip_patches": {},
|
|
"key_map": {"layer.weight": "alias"},
|
|
"output_strength": 1.0,
|
|
"clip_strength": 1.0,
|
|
},
|
|
tempfile.gettempdir(),
|
|
"../escape",
|
|
save_rank=0,
|
|
bake_strength=False,
|
|
)
|
|
|
|
with self.assertRaises(ValueError):
|
|
tuner_saver.save_tuner_data({"top_n": []}, tempfile.gettempdir(), "../escape")
|
|
|
|
def test_bridge_passthrough_returns_ui_payload(self):
|
|
tuner_data = {
|
|
"top_n": [{
|
|
"config": {
|
|
"optimization_mode": "global",
|
|
"merge_mode": "ties",
|
|
"merge_refinement": "none",
|
|
"sparsification": "disabled",
|
|
"sparsification_density": 0.7,
|
|
"dare_dampening": 0.0,
|
|
"auto_strength": "enabled",
|
|
},
|
|
"score_final": 0.91,
|
|
}],
|
|
"decision_smoothing": 0.25,
|
|
"auto_strength_floor": 0.95,
|
|
}
|
|
|
|
result = self.optimizer.execute_node(
|
|
self.model, [], 1.0,
|
|
tuner_data=tuner_data,
|
|
settings_source="from_autotuner",
|
|
)
|
|
|
|
self.assertIsInstance(result, dict)
|
|
self.assertIn("ui", result)
|
|
applied = json.loads(result["ui"]["applied_settings"][0])
|
|
self.assertEqual(applied["merge_mode"], "ties")
|
|
self.assertEqual(applied["auto_strength_floor"], 0.95)
|
|
|
|
def test_widget_order_keeps_upstream_workflow_compatibility(self):
|
|
optimizer_keys = list(lora_optimizer.LoRAOptimizer.INPUT_TYPES()["optional"].keys())
|
|
self.assertIn("settings_source", optimizer_keys)
|
|
self.assertIn("decision_smoothing", optimizer_keys)
|
|
self.assertLess(optimizer_keys.index("settings_source"), optimizer_keys.index("decision_smoothing"))
|
|
|
|
autotuner_keys = list(lora_optimizer.LoRAAutoTuner.INPUT_TYPES()["optional"].keys())
|
|
self.assertIn("output_mode", autotuner_keys)
|
|
self.assertIn("decision_smoothing", autotuner_keys)
|
|
self.assertLess(autotuner_keys.index("output_mode"), autotuner_keys.index("decision_smoothing"))
|
|
|
|
def test_optimizer_exposes_tuner_data_output_and_compatibility_analyzer_node(self):
|
|
self.assertEqual(
|
|
lora_optimizer.LoRAOptimizer.RETURN_TYPES,
|
|
("MODEL", "CLIP", "STRING", "TUNER_DATA", "LORA_DATA"),
|
|
)
|
|
self.assertIn("LoRACompatibilityAnalyzer", lora_optimizer.NODE_CLASS_MAPPINGS)
|
|
|
|
def test_autotune_resolve_tree_calls_auto_tune_for_subgroups(self):
|
|
"""_autotune_resolve_tree should call auto_tune for sub-groups with 2+ items."""
|
|
from lora_optimizer import LoRAAutoTuner, _parse_merge_formula
|
|
|
|
tuner = LoRAAutoTuner()
|
|
tree = _parse_merge_formula("(1+2)+3", 3)
|
|
|
|
# Build a minimal normalized stack with 3 fake LoRAs
|
|
fake_lora_a = {"key_a": torch.randn(4, 4)}
|
|
fake_lora_b = {"key_a": torch.randn(4, 4)}
|
|
fake_lora_c = {"key_c": torch.randn(4, 4)}
|
|
normalized_stack = [
|
|
{"name": "lora_a", "lora": fake_lora_a, "strength": 1.0,
|
|
"clip_strength": None, "metadata": {}},
|
|
{"name": "lora_b", "lora": fake_lora_b, "strength": 1.0,
|
|
"clip_strength": None, "metadata": {}},
|
|
{"name": "lora_c", "lora": fake_lora_c, "strength": 1.0,
|
|
"clip_strength": None, "metadata": {}},
|
|
]
|
|
|
|
# Track auto_tune calls
|
|
calls = []
|
|
|
|
def mock_auto_tune(model, lora_stack, output_strength, **kwargs):
|
|
calls.append({"n_loras": len(lora_stack), "names": [l["name"] for l in lora_stack]})
|
|
# Return a minimal 6-tuple with virtual LoRA patches
|
|
virtual_patches = {"key_a": ("diff", (torch.randn(4, 4),))}
|
|
lora_data = {"model_patches": virtual_patches, "clip_patches": {}}
|
|
return (model, None, "sub-report", "", None, lora_data)
|
|
|
|
tuner.auto_tune = mock_auto_tune
|
|
|
|
at_kwargs = {
|
|
"clip_strength_multiplier": 1.0,
|
|
"top_n": 3,
|
|
"normalize_keys": "disabled",
|
|
"scoring_svd": "disabled",
|
|
"scoring_device": "cpu",
|
|
"architecture_preset": "dit",
|
|
"auto_strength_floor": -1.0,
|
|
"decision_smoothing": 0.25,
|
|
"smooth_slerp_gate": False,
|
|
"vram_budget": 0.0,
|
|
"scoring_speed": "turbo",
|
|
"scoring_formula": "v2",
|
|
"diff_cache_mode": "disabled",
|
|
"diff_cache_ram_pct": 0.5,
|
|
}
|
|
|
|
resolved_stack, sub_reports = tuner._autotune_resolve_tree(
|
|
tree, normalized_stack, None, None, **at_kwargs)
|
|
|
|
# Should have called auto_tune once for the (1+2) sub-group
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertEqual(calls[0]["n_loras"], 2)
|
|
self.assertEqual(calls[0]["names"], ["lora_a", "lora_b"])
|
|
|
|
# Resolved stack should have 2 items: virtual LoRA + lora_c
|
|
self.assertEqual(len(resolved_stack), 2)
|
|
self.assertTrue(resolved_stack[0].get("_precomputed_diffs")) # virtual
|
|
self.assertEqual(resolved_stack[1]["name"], "lora_c")
|
|
self.assertEqual(len(sub_reports), 1)
|
|
|
|
def test_autotune_resolve_tree_passes_sub_merge_flags(self):
|
|
"""Sub-merge auto_tune calls should include _is_sub_merge and _suppress_pbar."""
|
|
from lora_optimizer import LoRAAutoTuner, _parse_merge_formula
|
|
|
|
tuner = LoRAAutoTuner()
|
|
tree = _parse_merge_formula("(1+2)+3", 3)
|
|
|
|
fake_lora_a = {"key_a": torch.randn(4, 4)}
|
|
fake_lora_b = {"key_a": torch.randn(4, 4)}
|
|
fake_lora_c = {"key_c": torch.randn(4, 4)}
|
|
normalized_stack = [
|
|
{"name": "lora_a", "lora": fake_lora_a, "strength": 1.0,
|
|
"clip_strength": None, "metadata": {}},
|
|
{"name": "lora_b", "lora": fake_lora_b, "strength": 1.0,
|
|
"clip_strength": None, "metadata": {}},
|
|
{"name": "lora_c", "lora": fake_lora_c, "strength": 1.0,
|
|
"clip_strength": None, "metadata": {}},
|
|
]
|
|
|
|
calls = []
|
|
def mock_auto_tune(model, lora_stack, output_strength, **kwargs):
|
|
calls.append(kwargs)
|
|
virtual_patches = {"key_a": ("diff", (torch.randn(4, 4),))}
|
|
lora_data = {"model_patches": virtual_patches, "clip_patches": {}}
|
|
return (model, None, "sub-report", "", None, lora_data)
|
|
|
|
tuner.auto_tune = mock_auto_tune
|
|
|
|
at_kwargs = {
|
|
"clip_strength_multiplier": 1.0,
|
|
"top_n": 3, "normalize_keys": "disabled", "scoring_svd": "disabled",
|
|
"scoring_device": "cpu", "architecture_preset": "dit",
|
|
"auto_strength_floor": -1.0, "decision_smoothing": 0.25,
|
|
"smooth_slerp_gate": False, "vram_budget": 0.0,
|
|
"scoring_speed": "turbo", "scoring_formula": "v2",
|
|
"diff_cache_mode": "disabled", "diff_cache_ram_pct": 0.5,
|
|
}
|
|
|
|
tuner._autotune_resolve_tree(tree, normalized_stack, None, None, **at_kwargs)
|
|
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertTrue(calls[0].get("_is_sub_merge", False))
|
|
self.assertTrue(calls[0].get("_suppress_pbar", False))
|
|
self.assertEqual(calls[0].get("cache_patches"), "disabled")
|
|
self.assertEqual(calls[0].get("community_cache"), "disabled")
|
|
self.assertEqual(calls[0].get("output_mode"), "merge")
|
|
|
|
def test_autotune_resolve_tree_single_item_passthrough(self):
|
|
"""Single-item sub-groups should pass through without calling auto_tune."""
|
|
from lora_optimizer import LoRAAutoTuner, _parse_merge_formula
|
|
|
|
tuner = LoRAAutoTuner()
|
|
tree = _parse_merge_formula("(1)+2", 2)
|
|
|
|
fake_lora_a = {"key_a": torch.randn(4, 4)}
|
|
fake_lora_b = {"key_b": torch.randn(4, 4)}
|
|
normalized_stack = [
|
|
{"name": "lora_a", "lora": fake_lora_a, "strength": 1.0,
|
|
"clip_strength": None, "metadata": {}},
|
|
{"name": "lora_b", "lora": fake_lora_b, "strength": 1.0,
|
|
"clip_strength": None, "metadata": {}},
|
|
]
|
|
|
|
calls = []
|
|
def mock_auto_tune(model, lora_stack, output_strength, **kwargs):
|
|
calls.append(True)
|
|
return (model, None, "", "", None, None)
|
|
|
|
tuner.auto_tune = mock_auto_tune
|
|
|
|
at_kwargs = {
|
|
"clip_strength_multiplier": 1.0,
|
|
"top_n": 3, "normalize_keys": "disabled", "scoring_svd": "disabled",
|
|
"scoring_device": "cpu", "architecture_preset": "dit",
|
|
"auto_strength_floor": -1.0, "decision_smoothing": 0.25,
|
|
"smooth_slerp_gate": False, "vram_budget": 0.0,
|
|
"scoring_speed": "turbo", "scoring_formula": "v2",
|
|
"diff_cache_mode": "disabled", "diff_cache_ram_pct": 0.5,
|
|
}
|
|
|
|
resolved_stack, sub_reports = tuner._autotune_resolve_tree(
|
|
tree, normalized_stack, None, None, **at_kwargs)
|
|
|
|
self.assertEqual(len(calls), 0)
|
|
self.assertEqual(len(resolved_stack), 2)
|
|
self.assertEqual(resolved_stack[0]["name"], "lora_a")
|
|
self.assertEqual(resolved_stack[1]["name"], "lora_b")
|
|
|
|
def test_autotune_resolve_tree_nested_groups(self):
|
|
"""Nested formula ((1+2)+3)+4 should resolve innermost first."""
|
|
from lora_optimizer import LoRAAutoTuner, _parse_merge_formula
|
|
|
|
tuner = LoRAAutoTuner()
|
|
tree = _parse_merge_formula("((1+2)+3)+4", 4)
|
|
|
|
normalized_stack = []
|
|
for i in range(4):
|
|
normalized_stack.append({
|
|
"name": f"lora_{i}", "lora": {f"key_{i}": torch.randn(4, 4)},
|
|
"strength": 1.0, "clip_strength": None, "metadata": {},
|
|
})
|
|
|
|
call_order = []
|
|
def mock_auto_tune(model, lora_stack, output_strength, **kwargs):
|
|
names = [l["name"] for l in lora_stack if not l.get("_precomputed_diffs")]
|
|
virtual_count = sum(1 for l in lora_stack if l.get("_precomputed_diffs"))
|
|
call_order.append({"names": names, "virtual_count": virtual_count,
|
|
"total": len(lora_stack)})
|
|
virtual_patches = {f"key_v{len(call_order)}": ("diff", (torch.randn(4, 4),))}
|
|
lora_data = {"model_patches": virtual_patches, "clip_patches": {}}
|
|
return (model, None, f"sub-report-{len(call_order)}", "", None, lora_data)
|
|
|
|
tuner.auto_tune = mock_auto_tune
|
|
|
|
at_kwargs = {
|
|
"clip_strength_multiplier": 1.0,
|
|
"top_n": 3, "normalize_keys": "disabled", "scoring_svd": "disabled",
|
|
"scoring_device": "cpu", "architecture_preset": "dit",
|
|
"auto_strength_floor": -1.0, "decision_smoothing": 0.25,
|
|
"smooth_slerp_gate": False, "vram_budget": 0.0,
|
|
"scoring_speed": "turbo", "scoring_formula": "v2",
|
|
"diff_cache_mode": "disabled", "diff_cache_ram_pct": 0.5,
|
|
}
|
|
|
|
resolved_stack, sub_reports = tuner._autotune_resolve_tree(
|
|
tree, normalized_stack, None, None, **at_kwargs)
|
|
|
|
# Two auto_tune calls: (1+2) first, then (virtual+3)
|
|
self.assertEqual(len(call_order), 2)
|
|
self.assertEqual(call_order[0]["names"], ["lora_0", "lora_1"])
|
|
self.assertEqual(call_order[0]["virtual_count"], 0)
|
|
self.assertEqual(call_order[1]["total"], 2)
|
|
self.assertEqual(call_order[1]["virtual_count"], 1)
|
|
|
|
# Final stack: virtual from ((1+2)+3) + lora_3
|
|
self.assertEqual(len(resolved_stack), 2)
|
|
self.assertTrue(resolved_stack[0].get("_precomputed_diffs"))
|
|
self.assertEqual(resolved_stack[1]["name"], "lora_3")
|
|
|
|
def test_auto_tune_with_formula_calls_autotune_resolve_tree(self):
|
|
"""auto_tune should detect formula and use _autotune_resolve_tree."""
|
|
tuner = lora_optimizer.LoRAAutoTuner()
|
|
|
|
# Build a lora_stack with formula metadata
|
|
fake_lora_a = {"key_a": torch.randn(4, 4)}
|
|
fake_lora_b = {"key_a": torch.randn(4, 4)}
|
|
fake_lora_c = {"key_c": torch.randn(4, 4)}
|
|
lora_stack = [
|
|
{"name": "lora_a", "lora": fake_lora_a, "strength": 1.0},
|
|
{"name": "lora_b", "lora": fake_lora_b, "strength": 1.0},
|
|
{"name": "lora_c", "lora": fake_lora_c, "strength": 1.0},
|
|
{"_merge_formula": "(1+2)+3"},
|
|
]
|
|
|
|
# Track _autotune_resolve_tree calls
|
|
resolve_calls = []
|
|
|
|
def mock_resolve(tree, normalized_stack, model, clip, **kwargs):
|
|
resolve_calls.append(tree)
|
|
# Return a flat stack (2 items) so auto_tune continues normally
|
|
return ([normalized_stack[0], normalized_stack[2]], [])
|
|
|
|
tuner._autotune_resolve_tree = mock_resolve
|
|
|
|
# Mock the rest of auto_tune to avoid needing real models
|
|
# We just need to verify _autotune_resolve_tree was called
|
|
try:
|
|
tuner.auto_tune(None, lora_stack, 1.0)
|
|
except Exception:
|
|
pass # Will fail later in pipeline — we only check the call happened
|
|
|
|
self.assertEqual(len(resolve_calls), 1)
|
|
self.assertEqual(resolve_calls[0]["type"], "group")
|
|
|
|
def test_build_stack_skips_disabled_slots(self):
|
|
from unittest import mock
|
|
node = lora_optimizer.LoRAStackDynamic()
|
|
with mock.patch.object(
|
|
lora_optimizer.LoRAStackDynamic, "_resolve_lora_name",
|
|
side_effect=lambda n: n,
|
|
):
|
|
result, = node.build_stack(
|
|
settings_visibility="simple",
|
|
input_mode="text",
|
|
lora_count=3,
|
|
lora_name_text_1="lora_a",
|
|
lora_name_text_2="lora_b",
|
|
lora_name_text_3="lora_c",
|
|
strength_1=1.0,
|
|
strength_2=0.8,
|
|
strength_3=0.5,
|
|
enabled_1=True,
|
|
enabled_2=False,
|
|
enabled_3=True,
|
|
)
|
|
names = [entry[0] for entry in result]
|
|
self.assertEqual(len(result), 2)
|
|
self.assertIn("lora_a", names)
|
|
self.assertIn("lora_c", names)
|
|
self.assertNotIn("lora_b", names)
|
|
|
|
def test_build_stack_absent_enabled_treated_as_enabled(self):
|
|
from unittest import mock
|
|
node = lora_optimizer.LoRAStackDynamic()
|
|
# No enabled_{i} keys passed — must behave as if all enabled
|
|
with mock.patch.object(
|
|
lora_optimizer.LoRAStackDynamic, "_resolve_lora_name",
|
|
side_effect=lambda n: n,
|
|
):
|
|
result, = node.build_stack(
|
|
settings_visibility="simple",
|
|
input_mode="text",
|
|
lora_count=2,
|
|
lora_name_text_1="lora_a",
|
|
lora_name_text_2="lora_b",
|
|
strength_1=1.0,
|
|
strength_2=0.8,
|
|
# No enabled_1 or enabled_2 passed
|
|
)
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_score_merge_result_baseline_matches_full(self):
|
|
"""Scoring multi-LoRA patches with single-LoRA baseline should match full scoring."""
|
|
LoRAAdapter = lora_optimizer.LoRAAdapter
|
|
single_patches = {}
|
|
multi_patches = {}
|
|
all_patches = {}
|
|
for i in range(10):
|
|
up = torch.randn(8, 4)
|
|
down = torch.randn(4, 16)
|
|
adapter = LoRAAdapter(set(), (up, down, 4.0, None, None, None))
|
|
key = f"key{i}"
|
|
all_patches[key] = adapter
|
|
if i < 5:
|
|
single_patches[key] = adapter
|
|
else:
|
|
multi_patches[key] = adapter
|
|
|
|
full = lora_optimizer._score_merge_result(all_patches, {}, compute_svd=False)
|
|
|
|
sl = lora_optimizer._score_merge_result(
|
|
single_patches, {}, compute_svd=False, _return_raw=True)
|
|
baseline = sl["_raw"]
|
|
combined = lora_optimizer._score_merge_result(
|
|
multi_patches, {}, compute_svd=False, _baseline=baseline)
|
|
|
|
self.assertAlmostEqual(full["composite_score"], combined["composite_score"], places=6)
|
|
self.assertAlmostEqual(full["norm_mean"], combined["norm_mean"], places=6)
|
|
self.assertAlmostEqual(full["norm_cv"], combined["norm_cv"], places=6)
|
|
self.assertAlmostEqual(full["sparsity_mean"], combined["sparsity_mean"], places=6)
|
|
self.assertAlmostEqual(full["norm_energy_sq"], combined["norm_energy_sq"], places=4)
|
|
|
|
|
|
@unittest.skipIf(torch is None, "torch is not installed in this environment")
|
|
class LoRASettingsNodeTests(unittest.TestCase):
|
|
"""Tests for LoRAMergeSettings, LoRAOptimizerSettings and LoRAAutoTunerSettings nodes."""
|
|
|
|
def _build_defaults(self, inputs):
|
|
"""Extract default values from INPUT_TYPES required spec."""
|
|
defaults = {}
|
|
for key, spec in inputs["required"].items():
|
|
if isinstance(spec[0], list):
|
|
defaults[key] = spec[1].get("default", spec[0][0])
|
|
elif spec[0] == "FLOAT":
|
|
defaults[key] = spec[1]["default"]
|
|
elif spec[0] == "INT":
|
|
defaults[key] = spec[1]["default"]
|
|
elif spec[0] == "BOOLEAN":
|
|
defaults[key] = spec[1]["default"]
|
|
return defaults
|
|
|
|
def test_merge_settings_build_returns_dict(self):
|
|
node = lora_optimizer.LoRAMergeSettings()
|
|
inputs = lora_optimizer.LoRAMergeSettings.INPUT_TYPES()
|
|
defaults = self._build_defaults(inputs)
|
|
result = node.build_settings(**defaults)
|
|
self.assertIsInstance(result, tuple)
|
|
self.assertEqual(len(result), 1)
|
|
settings = result[0]
|
|
expected_keys = {"normalize_keys", "architecture_preset",
|
|
"auto_strength_floor", "decision_smoothing",
|
|
"smooth_slerp_gate", "vram_budget", "cache_patches"}
|
|
self.assertEqual(set(settings.keys()), expected_keys)
|
|
self.assertEqual(settings["normalize_keys"], "enabled")
|
|
self.assertEqual(settings["architecture_preset"], "auto")
|
|
self.assertAlmostEqual(settings["auto_strength_floor"], -1.0)
|
|
self.assertAlmostEqual(settings["decision_smoothing"], 0.25)
|
|
self.assertFalse(settings["smooth_slerp_gate"])
|
|
self.assertAlmostEqual(settings["vram_budget"], 0.0)
|
|
self.assertEqual(settings["cache_patches"], "enabled")
|
|
|
|
def test_optimizer_settings_build_returns_advanced_mode(self):
|
|
node = lora_optimizer.LoRAOptimizerSettings()
|
|
inputs = lora_optimizer.LoRAOptimizerSettings.INPUT_TYPES()
|
|
defaults = self._build_defaults(inputs)
|
|
result = node.build_settings(**defaults)
|
|
self.assertIsInstance(result, tuple)
|
|
self.assertEqual(len(result), 1)
|
|
settings = result[0]
|
|
self.assertEqual(settings["mode"], "advanced")
|
|
self.assertEqual(settings["auto_strength"], "enabled")
|
|
self.assertEqual(settings["optimization_mode"], "per_prefix")
|
|
self.assertEqual(settings["sparsification"], "disabled")
|
|
self.assertAlmostEqual(settings["sparsification_density"], 0.7)
|
|
self.assertEqual(settings["merge_strategy_override"], "")
|
|
# Common settings should use defaults when merge_settings not connected
|
|
self.assertEqual(settings["normalize_keys"], "enabled")
|
|
self.assertEqual(settings["architecture_preset"], "auto")
|
|
|
|
def test_optimizer_settings_with_strategy_override(self):
|
|
node = lora_optimizer.LoRAOptimizerSettings()
|
|
inputs = lora_optimizer.LoRAOptimizerSettings.INPUT_TYPES()
|
|
defaults = self._build_defaults(inputs)
|
|
defaults["merge_strategy_override"] = "slerp"
|
|
result = node.build_settings(**defaults)
|
|
self.assertEqual(result[0]["merge_strategy_override"], "slerp")
|
|
|
|
def test_autotuner_settings_build_returns_autotuner_mode(self):
|
|
node = lora_optimizer.LoRAAutoTunerSettings()
|
|
inputs = lora_optimizer.LoRAAutoTunerSettings.INPUT_TYPES()
|
|
defaults = self._build_defaults(inputs)
|
|
result = node.build_settings(**defaults)
|
|
self.assertIsInstance(result, tuple)
|
|
self.assertEqual(len(result), 1)
|
|
settings = result[0]
|
|
self.assertEqual(settings["mode"], "autotuner")
|
|
self.assertEqual(settings["top_n"], 3)
|
|
self.assertEqual(settings["scoring_speed"], "turbo")
|
|
self.assertEqual(settings["output_mode"], "merge")
|
|
self.assertFalse(settings["smooth_slerp_gate"])
|
|
self.assertIsNone(settings["evaluator"])
|
|
# Common settings should use defaults when merge_settings not connected
|
|
self.assertEqual(settings["normalize_keys"], "enabled")
|
|
self.assertEqual(settings["cache_patches"], "enabled")
|
|
|
|
def test_autotuner_settings_with_evaluator(self):
|
|
node = lora_optimizer.LoRAAutoTunerSettings()
|
|
inputs = lora_optimizer.LoRAAutoTunerSettings.INPUT_TYPES()
|
|
defaults = self._build_defaults(inputs)
|
|
evaluator = {"type": "python", "code": "return 0.5"}
|
|
defaults["evaluator"] = evaluator
|
|
result = node.build_settings(**defaults)
|
|
self.assertEqual(result[0]["evaluator"], evaluator)
|
|
|
|
def test_mode_settings_merge_base_settings(self):
|
|
"""Both mode nodes correctly merge base settings from LoRAMergeSettings."""
|
|
custom_ms = {
|
|
"normalize_keys": "disabled",
|
|
"architecture_preset": "dit",
|
|
"auto_strength_floor": 0.5,
|
|
"decision_smoothing": 0.8,
|
|
"smooth_slerp_gate": True,
|
|
"vram_budget": 0.3,
|
|
"cache_patches": "disabled",
|
|
}
|
|
|
|
# Test LoRAOptimizerSettings with custom merge_settings
|
|
opt_node = lora_optimizer.LoRAOptimizerSettings()
|
|
opt_inputs = lora_optimizer.LoRAOptimizerSettings.INPUT_TYPES()
|
|
opt_defaults = self._build_defaults(opt_inputs)
|
|
opt_defaults["merge_settings"] = custom_ms
|
|
opt_result = opt_node.build_settings(**opt_defaults)[0]
|
|
for key, val in custom_ms.items():
|
|
self.assertEqual(opt_result[key], val,
|
|
f"OptimizerSettings: {key} should be {val}, got {opt_result[key]}")
|
|
|
|
# Test LoRAAutoTunerSettings with custom merge_settings
|
|
at_node = lora_optimizer.LoRAAutoTunerSettings()
|
|
at_inputs = lora_optimizer.LoRAAutoTunerSettings.INPUT_TYPES()
|
|
at_defaults = self._build_defaults(at_inputs)
|
|
at_defaults["merge_settings"] = custom_ms
|
|
at_result = at_node.build_settings(**at_defaults)[0]
|
|
for key, val in custom_ms.items():
|
|
self.assertEqual(at_result[key], val,
|
|
f"AutoTunerSettings: {key} should be {val}, got {at_result[key]}")
|
|
|
|
def test_settings_nodes_registered_in_mappings(self):
|
|
self.assertIn("LoRAMergeSettings", lora_optimizer.NODE_CLASS_MAPPINGS)
|
|
self.assertIn("LoRAOptimizerSettings", lora_optimizer.NODE_CLASS_MAPPINGS)
|
|
self.assertIn("LoRAAutoTunerSettings", lora_optimizer.NODE_CLASS_MAPPINGS)
|
|
self.assertEqual(
|
|
lora_optimizer.NODE_DISPLAY_NAME_MAPPINGS["LoRAMergeSettings"],
|
|
"LoRA Merge Settings",
|
|
)
|
|
self.assertEqual(
|
|
lora_optimizer.NODE_DISPLAY_NAME_MAPPINGS["LoRAOptimizerSettings"],
|
|
"LoRA Optimizer Settings",
|
|
)
|
|
self.assertEqual(
|
|
lora_optimizer.NODE_DISPLAY_NAME_MAPPINGS["LoRAAutoTunerSettings"],
|
|
"LoRA AutoTuner Settings",
|
|
)
|
|
|
|
def test_settings_nodes_return_types(self):
|
|
self.assertEqual(lora_optimizer.LoRAMergeSettings.RETURN_TYPES, ("MERGE_SETTINGS",))
|
|
self.assertEqual(lora_optimizer.LoRAOptimizerSettings.RETURN_TYPES, ("OPTIMIZER_SETTINGS",))
|
|
self.assertEqual(lora_optimizer.LoRAAutoTunerSettings.RETURN_TYPES, ("OPTIMIZER_SETTINGS",))
|
|
|
|
def test_simple_node_accepts_settings_input(self):
|
|
inputs = lora_optimizer.LoRAOptimizerSimple.INPUT_TYPES()
|
|
self.assertIn("settings", inputs["optional"])
|
|
self.assertEqual(inputs["optional"]["settings"][0], "OPTIMIZER_SETTINGS")
|
|
|
|
def test_simple_is_changed_includes_settings_hash(self):
|
|
settings = {"mode": "advanced", "auto_strength": "enabled"}
|
|
result_with = lora_optimizer.LoRAOptimizerSimple.IS_CHANGED(
|
|
None, None, 1.0, settings=settings)
|
|
result_without = lora_optimizer.LoRAOptimizerSimple.IS_CHANGED(
|
|
None, None, 1.0)
|
|
self.assertIn("|settings=", result_with)
|
|
self.assertNotIn("|settings=", result_without)
|
|
|
|
def test_simple_is_changed_different_settings_produce_different_hashes(self):
|
|
settings_a = {"mode": "advanced", "auto_strength": "enabled"}
|
|
settings_b = {"mode": "advanced", "auto_strength": "disabled"}
|
|
result_a = lora_optimizer.LoRAOptimizerSimple.IS_CHANGED(
|
|
None, None, 1.0, settings=settings_a)
|
|
result_b = lora_optimizer.LoRAOptimizerSimple.IS_CHANGED(
|
|
None, None, 1.0, settings=settings_b)
|
|
self.assertNotEqual(result_a, result_b)
|
|
|
|
def test_optimizer_settings_defaults_match_simple_defaults(self):
|
|
"""Verify LoRAOptimizerSettings + LoRAMergeSettings defaults align with _SIMPLE_DEFAULTS."""
|
|
opt_inputs = lora_optimizer.LoRAOptimizerSettings.INPUT_TYPES()
|
|
merge_inputs = lora_optimizer.LoRAMergeSettings.INPUT_TYPES()
|
|
simple = lora_optimizer.LoRAOptimizerSimple._SIMPLE_DEFAULTS
|
|
# Keys on LoRAOptimizerSettings
|
|
for key in ["auto_strength", "optimization_mode", "sparsification",
|
|
"merge_refinement", "strategy_set", "patch_compression",
|
|
"svd_device", "free_vram_between_passes"]:
|
|
spec = opt_inputs["required"][key]
|
|
if isinstance(spec[0], list):
|
|
default = spec[1].get("default", spec[0][0])
|
|
else:
|
|
default = spec[1]["default"]
|
|
self.assertEqual(default, simple[key],
|
|
f"Default mismatch for {key}: settings={default}, simple={simple[key]}")
|
|
# Keys on LoRAMergeSettings
|
|
for key in ["normalize_keys", "architecture_preset", "cache_patches",
|
|
"smooth_slerp_gate"]:
|
|
spec = merge_inputs["required"][key]
|
|
if isinstance(spec[0], list):
|
|
default = spec[1].get("default", spec[0][0])
|
|
else:
|
|
default = spec[1]["default"]
|
|
self.assertEqual(default, simple[key],
|
|
f"Default mismatch for {key}: merge_settings={default}, simple={simple[key]}")
|
|
|
|
def test_merge_settings_defaults_match_input_types(self):
|
|
"""Verify _DEFAULTS dict stays in sync with INPUT_TYPES defaults."""
|
|
inputs = lora_optimizer.LoRAMergeSettings.INPUT_TYPES()
|
|
defaults_dict = lora_optimizer.LoRAMergeSettings._DEFAULTS
|
|
for key, spec in inputs["required"].items():
|
|
if isinstance(spec[0], list):
|
|
input_default = spec[1].get("default", spec[0][0])
|
|
else:
|
|
input_default = spec[1]["default"]
|
|
self.assertEqual(defaults_dict[key], input_default,
|
|
f"_DEFAULTS[{key}]={defaults_dict[key]} != INPUT_TYPES default={input_default}")
|
|
self.assertEqual(set(defaults_dict.keys()), set(inputs["required"].keys()),
|
|
"_DEFAULTS keys don't match INPUT_TYPES keys")
|
|
|
|
|
|
# --- Merge formula parser tests ---
|
|
|
|
def test_parse_merge_formula_simple(self):
|
|
"""Simple flat formula parses to group of leaves."""
|
|
tree = lora_optimizer._parse_merge_formula("1 + 2 + 3", 3)
|
|
self.assertEqual(tree["type"], "group")
|
|
self.assertEqual(len(tree["children"]), 3)
|
|
for i, child in enumerate(tree["children"]):
|
|
self.assertEqual(child["type"], "leaf")
|
|
self.assertEqual(child["index"], i)
|
|
|
|
def test_parse_merge_formula_nested(self):
|
|
"""Nested formula parses to tree with sub-group."""
|
|
tree = lora_optimizer._parse_merge_formula("(1+2) + 3", 3)
|
|
self.assertEqual(tree["type"], "group")
|
|
self.assertEqual(len(tree["children"]), 2)
|
|
sub = tree["children"][0]
|
|
self.assertEqual(sub["type"], "group")
|
|
self.assertEqual(len(sub["children"]), 2)
|
|
leaf3 = tree["children"][1]
|
|
self.assertEqual(leaf3["type"], "leaf")
|
|
self.assertEqual(leaf3["index"], 2)
|
|
|
|
def test_parse_merge_formula_weights(self):
|
|
"""Weights are parsed from :N.N suffix."""
|
|
tree = lora_optimizer._parse_merge_formula("(1+2):0.6 + 3:0.4", 3)
|
|
self.assertAlmostEqual(tree["children"][0]["weight"], 0.6)
|
|
self.assertAlmostEqual(tree["children"][1]["weight"], 0.4)
|
|
|
|
def test_parse_merge_formula_deep_nesting(self):
|
|
"""Deep nesting: ((1+2)+3) + 4."""
|
|
tree = lora_optimizer._parse_merge_formula("((1+2)+3) + 4", 4)
|
|
self.assertEqual(tree["type"], "group")
|
|
self.assertEqual(len(tree["children"]), 2)
|
|
inner = tree["children"][0]
|
|
self.assertEqual(inner["type"], "group")
|
|
self.assertEqual(len(inner["children"]), 2)
|
|
innermost = inner["children"][0]
|
|
self.assertEqual(innermost["type"], "group")
|
|
self.assertEqual(len(innermost["children"]), 2)
|
|
|
|
def test_parse_merge_formula_single_item(self):
|
|
"""Single item is valid."""
|
|
tree = lora_optimizer._parse_merge_formula("1", 1)
|
|
self.assertEqual(tree["type"], "leaf")
|
|
self.assertEqual(tree["index"], 0)
|
|
|
|
def test_parse_merge_formula_out_of_range(self):
|
|
"""Out of range index raises ValueError."""
|
|
with self.assertRaises(ValueError):
|
|
lora_optimizer._parse_merge_formula("1 + 5", 3)
|
|
|
|
def test_parse_merge_formula_malformed(self):
|
|
"""Malformed formula raises ValueError."""
|
|
with self.assertRaises(ValueError):
|
|
lora_optimizer._parse_merge_formula("((1+2", 3)
|
|
|
|
def test_parse_merge_formula_empty(self):
|
|
"""Empty/whitespace formula raises ValueError."""
|
|
with self.assertRaises(ValueError):
|
|
lora_optimizer._parse_merge_formula("", 3)
|
|
with self.assertRaises(ValueError):
|
|
lora_optimizer._parse_merge_formula(" ", 3)
|
|
|
|
|
|
def test_merge_formula_node_registered(self):
|
|
"""LoRAMergeFormula is registered in NODE_CLASS_MAPPINGS."""
|
|
self.assertIn("LoRAMergeFormula", lora_optimizer.NODE_CLASS_MAPPINGS)
|
|
self.assertIn("LoRAMergeFormula", lora_optimizer.NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
def test_merge_formula_node_passthrough(self):
|
|
"""LoRAMergeFormula passes stack through with formula metadata."""
|
|
node = lora_optimizer.LoRAMergeFormula()
|
|
stack = [{"name": "a", "lora": {}, "strength": 1.0}]
|
|
result = node.apply_formula(stack, "(1)")
|
|
self.assertIsInstance(result, tuple)
|
|
output_stack = result[0]
|
|
has_formula = any(isinstance(item, dict) and "_merge_formula" in item for item in output_stack)
|
|
self.assertTrue(has_formula)
|
|
|
|
def test_merge_formula_node_validates(self):
|
|
"""LoRAMergeFormula validates formula syntax — invalid returns stack without formula."""
|
|
node = lora_optimizer.LoRAMergeFormula()
|
|
stack = [{"name": "a", "lora": {}, "strength": 1.0}]
|
|
result = node.apply_formula(stack, "(1+2)") # only 1 LoRA — out of range
|
|
output_stack = result[0]
|
|
self.assertIsInstance(output_stack, list)
|
|
# Should NOT have formula metadata since validation failed
|
|
has_formula = any(isinstance(item, dict) and "_merge_formula" in item for item in output_stack)
|
|
self.assertFalse(has_formula)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Merge formula tree executor + optimize_merge integration
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_normalize_stack_filters_formula_metadata(self):
|
|
"""_normalize_stack filters out formula metadata entries."""
|
|
stack = [
|
|
{"name": "a", "lora": {}, "strength": 1.0},
|
|
{"_merge_formula": "(1+2)"},
|
|
{"name": "b", "lora": {}, "strength": 1.0},
|
|
]
|
|
opt = lora_optimizer.LoRAOptimizer()
|
|
result = opt._normalize_stack(stack)
|
|
self.assertEqual(len(result), 2)
|
|
names = [item["name"] for item in result]
|
|
self.assertEqual(names, ["a", "b"])
|
|
|
|
def test_optimize_merge_extracts_formula_metadata(self):
|
|
"""optimize_merge strips formula metadata before normalization."""
|
|
stack = [
|
|
{"name": "a", "lora": {"key1": ("diff", (torch.randn(4, 4),))}, "strength": 1.0},
|
|
{"name": "b", "lora": {"key1": ("diff", (torch.randn(4, 4),))}, "strength": 1.0},
|
|
{"_merge_formula": "1 + 2"},
|
|
]
|
|
opt = lora_optimizer.LoRAOptimizer()
|
|
# model=None → optimize_merge returns early with a report (no model to patch)
|
|
result = opt.optimize_merge(None, stack, 1.0)
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(len(result), 5) # 5-tuple
|
|
|
|
def test_optimize_merge_invalid_formula_falls_back(self):
|
|
"""Invalid formula logs a warning and falls back to flat merge."""
|
|
stack = [
|
|
{"name": "a", "lora": {"key1": ("diff", (torch.randn(4, 4),))}, "strength": 1.0},
|
|
{"name": "b", "lora": {"key1": ("diff", (torch.randn(4, 4),))}, "strength": 1.0},
|
|
{"_merge_formula": "((1+2"}, # malformed
|
|
]
|
|
opt = lora_optimizer.LoRAOptimizer()
|
|
result = opt.optimize_merge(None, stack, 1.0)
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(len(result), 5)
|
|
|
|
def test_model_to_virtual_lora_label(self):
|
|
"""_model_to_virtual_lora produces correct tree labels."""
|
|
tree_node = {
|
|
"type": "group",
|
|
"weight": None,
|
|
"children": [
|
|
{"type": "leaf", "index": 0, "weight": None},
|
|
{"type": "leaf", "index": 1, "weight": None},
|
|
],
|
|
}
|
|
virtual = lora_optimizer.LoRAOptimizer._model_to_virtual_lora(
|
|
{}, {}, tree_node)
|
|
self.assertEqual(virtual["name"], "(1+2)")
|
|
self.assertEqual(virtual["strength"], 1.0)
|
|
self.assertIsInstance(virtual["lora"], dict)
|
|
self.assertTrue(virtual["_precomputed_diffs"])
|
|
|
|
def test_resolve_tree_to_stack_flat(self):
|
|
"""_resolve_tree_to_stack resolves a flat group to the original items."""
|
|
opt = lora_optimizer.LoRAOptimizer()
|
|
normalized = [
|
|
{"name": "a", "lora": {}, "strength": 0.8, "clip_strength": None,
|
|
"conflict_mode": "all", "key_filter": "all", "metadata": {}},
|
|
{"name": "b", "lora": {}, "strength": 0.6, "clip_strength": None,
|
|
"conflict_mode": "all", "key_filter": "all", "metadata": {}},
|
|
]
|
|
tree = {
|
|
"type": "group",
|
|
"weight": None,
|
|
"children": [
|
|
{"type": "leaf", "index": 0, "weight": None},
|
|
{"type": "leaf", "index": 1, "weight": 0.5},
|
|
],
|
|
}
|
|
resolved, reports = opt._resolve_tree_to_stack(tree, normalized, None, None)
|
|
self.assertEqual(len(resolved), 2)
|
|
self.assertEqual(resolved[0]["strength"], 0.8) # unchanged
|
|
self.assertEqual(resolved[1]["strength"], 0.5) # overridden by weight
|
|
self.assertEqual(reports, [])
|
|
|
|
|
|
@unittest.skipIf(torch is None, "torch not available")
|
|
class TestExtractLoRAFromDelta(unittest.TestCase):
|
|
"""Tests for _extract_lora_svd() helper used by LoRAExtractFromModel."""
|
|
|
|
def _make_delta(self, rows, cols, rank):
|
|
"""Create a clean rank-r delta matrix using orthonormal U and V."""
|
|
U, _ = torch.linalg.qr(torch.randn(rows, rank)) # orthonormal (rows, rank)
|
|
V, _ = torch.linalg.qr(torch.randn(cols, rank)) # orthonormal (cols, rank)
|
|
S = torch.rand(rank) + 0.5 # positive singular values, not too small
|
|
return (U * S.unsqueeze(0)) @ V.T # (rows, cols)
|
|
|
|
def test_fixed_rank_output_shape(self):
|
|
"""Fixed mode: output lora_up/down have correct shapes."""
|
|
mod = _load_module()
|
|
delta = self._make_delta(64, 32, 8)
|
|
up, down, alpha = mod._extract_lora_svd(delta, rank=4, rank_mode="fixed", energy_threshold=0.99)
|
|
self.assertEqual(up.shape, (64, 4))
|
|
self.assertEqual(down.shape, (4, 32))
|
|
self.assertEqual(alpha, 4.0)
|
|
|
|
def test_auto_rank_energy_retained(self):
|
|
"""Auto mode: retained energy >= threshold."""
|
|
mod = _load_module()
|
|
delta = self._make_delta(64, 32, 16)
|
|
up, down, alpha = mod._extract_lora_svd(delta, rank=32, rank_mode="auto", energy_threshold=0.95)
|
|
# Reconstruct and check energy
|
|
reconstructed = up @ down
|
|
original_energy = (delta ** 2).sum().item()
|
|
reconstructed_energy = (reconstructed ** 2).sum().item()
|
|
self.assertGreaterEqual(reconstructed_energy / original_energy, 0.93) # allow small numerical error
|
|
|
|
def test_near_zero_delta_returns_none(self):
|
|
"""Near-zero delta (unaffected layer) returns None."""
|
|
mod = _load_module()
|
|
delta = torch.zeros(64, 32)
|
|
result = mod._extract_lora_svd(delta, rank=4, rank_mode="fixed", energy_threshold=0.99)
|
|
self.assertIsNone(result)
|
|
|
|
def test_singular_value_floor_applied(self):
|
|
"""Noise-only singular values below floor are excluded even in fixed mode."""
|
|
mod = _load_module()
|
|
# Create a clean rank-2 signal with tiny noise
|
|
signal = self._make_delta(32, 32, 2)
|
|
noise = torch.randn(32, 32) * 1e-6
|
|
delta = signal + noise
|
|
up, down, alpha = mod._extract_lora_svd(delta, rank=16, rank_mode="fixed", energy_threshold=0.99)
|
|
# Floor should cut noise singular values — effective rank must be ≤ 2
|
|
self.assertLessEqual(alpha, 2.0)
|
|
|
|
|
|
def test_conv_layer_reshaped_to_2d(self):
|
|
"""4D conv delta is reshaped to 2D; lora_down has flat spatial dimension."""
|
|
mod = _load_module()
|
|
# Simulate a conv2d weight delta: (C_out=64, C_in=32, kH=3, kW=3)
|
|
delta = torch.randn(64, 32, 3, 3)
|
|
up, down, alpha = mod._extract_lora_svd(delta, rank=4, rank_mode="fixed", energy_threshold=0.99)
|
|
self.assertEqual(up.shape[0], 64) # rows = C_out
|
|
self.assertEqual(down.shape[1], 32 * 3 * 3) # cols = C_in * kH * kW
|
|
self.assertEqual(up.shape[1], down.shape[0]) # inner dim consistent
|
|
|
|
def test_1d_delta_returns_none(self):
|
|
"""1D bias delta returns None (SVD requires 2D input)."""
|
|
mod = _load_module()
|
|
delta = torch.randn(64) # bias vector
|
|
result = mod._extract_lora_svd(delta, rank=4, rank_mode="fixed", energy_threshold=0.99)
|
|
self.assertIsNone(result)
|
|
|
|
|
|
@unittest.skipIf(torch is None, "torch is not installed")
|
|
class AnalysisCacheTests(unittest.TestCase):
|
|
def test_names_only_hash_excludes_strength(self):
|
|
"""Same LoRA files at different strengths produce the same hash."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path_a = os.path.join(tmpdir, "lora_a.safetensors")
|
|
path_b = os.path.join(tmpdir, "lora_b.safetensors")
|
|
open(path_a, "wb").close()
|
|
open(path_b, "wb").close()
|
|
|
|
with mock.patch("lora_optimizer.folder_paths.get_full_path",
|
|
side_effect=lambda _k, n: os.path.join(tmpdir, n)):
|
|
stack_s1 = [
|
|
{"name": "lora_a.safetensors", "strength": 0.5},
|
|
{"name": "lora_b.safetensors", "strength": 1.0},
|
|
]
|
|
stack_s2 = [
|
|
{"name": "lora_a.safetensors", "strength": 1.5},
|
|
{"name": "lora_b.safetensors", "strength": 0.3},
|
|
]
|
|
h1, signs1 = lora_optimizer.LoRAAutoTuner._compute_names_only_hash(stack_s1)
|
|
h2, signs2 = lora_optimizer.LoRAAutoTuner._compute_names_only_hash(stack_s2)
|
|
self.assertEqual(h1, h2)
|
|
self.assertEqual(signs1, {0: 1, 1: 1})
|
|
self.assertEqual(signs2, {0: 1, 1: 1})
|
|
|
|
def test_names_only_hash_captures_sign(self):
|
|
"""Negative strength is captured in the returned signs dict."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path_a = os.path.join(tmpdir, "lora_a.safetensors")
|
|
open(path_a, "wb").close()
|
|
with mock.patch("lora_optimizer.folder_paths.get_full_path",
|
|
side_effect=lambda _k, n: os.path.join(tmpdir, n)):
|
|
stack_pos = [{"name": "lora_a.safetensors", "strength": 1.0}]
|
|
stack_neg = [{"name": "lora_a.safetensors", "strength": -1.0}]
|
|
_, signs_pos = lora_optimizer.LoRAAutoTuner._compute_names_only_hash(stack_pos)
|
|
_, signs_neg = lora_optimizer.LoRAAutoTuner._compute_names_only_hash(stack_neg)
|
|
self.assertEqual(signs_pos[0], 1)
|
|
self.assertEqual(signs_neg[0], -1)
|
|
|
|
def test_analysis_cache_roundtrip(self):
|
|
"""Save and load analysis cache; content survives round-trip."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
per_prefix = {
|
|
"prefix_a": {
|
|
"pair_conflicts": {"0,1": {"overlap": 100, "conflict": 30,
|
|
"dot": 0.5, "norm_a_sq": 1.0,
|
|
"norm_b_sq": 1.0,
|
|
"weighted_total": 0.8,
|
|
"weighted_conflict": 0.2,
|
|
"expected_conflict": 0.15,
|
|
"excess_conflict": 0.05,
|
|
"subspace_overlap": 0.3,
|
|
"subspace_weight": 1.0}},
|
|
"per_lora_norm_sq": {"0": 1.5, "1": 0.8},
|
|
"magnitude_samples_unscaled": {"0": [0.1, 0.2], "1": [0.3]},
|
|
"ranks": {"0": 16, "1": 32},
|
|
"target_key": "model.layer.weight",
|
|
"is_clip": False,
|
|
"raw_n": 2,
|
|
"skip_count": 0,
|
|
"strength_signs": {"0": 1, "1": 1},
|
|
}
|
|
}
|
|
source_loras = [{"name": "a.safetensors", "mtime": 1.0, "size": 100}]
|
|
lora_optimizer.LoRAAutoTuner._analysis_cache_save(
|
|
"abc123", per_prefix, source_loras)
|
|
|
|
loaded = lora_optimizer.LoRAAutoTuner._analysis_cache_load("abc123")
|
|
self.assertIsNotNone(loaded)
|
|
self.assertIn("prefix_a", loaded)
|
|
self.assertEqual(loaded["prefix_a"]["per_lora_norm_sq"]["0"], 1.5)
|
|
|
|
def test_analysis_cache_load_missing_returns_none(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
self.assertIsNone(
|
|
lora_optimizer.LoRAAutoTuner._analysis_cache_load("nonexistent"))
|
|
|
|
def test_analysis_cache_load_stale_algo_version_returns_none(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
stale_hash = "staletest1"
|
|
path = os.path.join(tmpdir, f"{stale_hash}.analysis.json")
|
|
with open(path, "w") as f:
|
|
json.dump({"analysis_version": 1,
|
|
"algo_version": "0.0.0",
|
|
"per_prefix": {"prefix_a": {}}}, f)
|
|
result = lora_optimizer.LoRAAutoTuner._analysis_cache_load(stale_hash)
|
|
self.assertIsNone(result)
|
|
|
|
def test_analysis_partial_path_uses_partial_suffix(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
path = lora_optimizer.LoRAAutoTuner._analysis_partial_path("abc123")
|
|
self.assertTrue(path.endswith("abc123.analysis.partial.json"))
|
|
|
|
def test_analysis_partial_load_missing_returns_none(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
self.assertIsNone(
|
|
lora_optimizer.LoRAAutoTuner._analysis_partial_load("nonexistent"))
|
|
|
|
def test_analysis_partial_load_stale_algo_version_returns_none(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
path = os.path.join(tmpdir, "stale.analysis.partial.json")
|
|
with open(path, "w") as f:
|
|
json.dump({"algo_version": "0.0.0", "per_prefix": {}}, f)
|
|
self.assertIsNone(
|
|
lora_optimizer.LoRAAutoTuner._analysis_partial_load("stale"))
|
|
|
|
def test_analysis_partial_roundtrip(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
per_prefix = {"prefix_a": {"ranks": {"0": 16}, "is_clip": False}}
|
|
source_loras = [{"name": "a.safetensors"}]
|
|
lora_optimizer.LoRAAutoTuner._analysis_partial_save(
|
|
"abc123", per_prefix, source_loras)
|
|
loaded = lora_optimizer.LoRAAutoTuner._analysis_partial_load("abc123")
|
|
self.assertIsNotNone(loaded)
|
|
self.assertIn("prefix_a", loaded)
|
|
|
|
def test_analysis_partial_save_is_atomic(self):
|
|
"""Save uses tmp+replace so a crash mid-write doesn't corrupt the file."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
lora_optimizer.LoRAAutoTuner._analysis_partial_save(
|
|
"abc123", {"p": {}}, [])
|
|
partial_path = lora_optimizer.LoRAAutoTuner._analysis_partial_path("abc123")
|
|
tmp_path = partial_path + ".tmp"
|
|
self.assertTrue(os.path.exists(partial_path))
|
|
self.assertFalse(os.path.exists(tmp_path))
|
|
|
|
def test_analysis_partial_delete_removes_file(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
lora_optimizer.LoRAAutoTuner._analysis_partial_save(
|
|
"abc123", {}, [])
|
|
path = lora_optimizer.LoRAAutoTuner._analysis_partial_path("abc123")
|
|
self.assertTrue(os.path.exists(path))
|
|
lora_optimizer.LoRAAutoTuner._analysis_partial_delete("abc123")
|
|
self.assertFalse(os.path.exists(path))
|
|
|
|
def test_analysis_partial_delete_silent_on_missing(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
# Should not raise
|
|
lora_optimizer.LoRAAutoTuner._analysis_partial_delete("nonexistent")
|
|
|
|
def test_extract_for_analysis_cache_strips_strength(self):
|
|
"""Extracted data has unscaled magnitude samples; tuple keys become strings."""
|
|
partial_stats = [(0, 16, 1.5, 2.25), (1, 32, 0.9, 0.81)]
|
|
pair_conflicts = {
|
|
(0, 1): {"overlap": 50, "conflict": 10, "dot": 0.4,
|
|
"norm_a_sq": 2.25, "norm_b_sq": 0.81,
|
|
"weighted_total": 0.6, "weighted_conflict": 0.1,
|
|
"expected_conflict": 0.12, "excess_conflict": 0.0,
|
|
"subspace_overlap": 0.2, "subspace_weight": 1.35}
|
|
}
|
|
magnitude_samples = [
|
|
torch.tensor([1.5, 3.0]), # LoRA 0: scaled by abs(strength=1.5)
|
|
torch.tensor([0.9, 1.8]), # LoRA 1: scaled by abs(strength=0.9)
|
|
]
|
|
per_lora_norm_sq = {0: 2.25, 1: 0.81}
|
|
result = (
|
|
"prefix_x", partial_stats, pair_conflicts,
|
|
magnitude_samples, ("layer.weight", False),
|
|
0, 2, per_lora_norm_sq,
|
|
)
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 1.5, "clip_strength": None},
|
|
{"name": "b.safetensors", "strength": 0.9, "clip_strength": None},
|
|
]
|
|
extracted = lora_optimizer.LoRAOptimizer._extract_for_analysis_cache(
|
|
result, active_loras)
|
|
|
|
# Pair key must be a string
|
|
self.assertIn("0,1", extracted["pair_conflicts"])
|
|
# Samples unscaled: divide by abs(strength)
|
|
self.assertAlmostEqual(extracted["magnitude_samples_unscaled"]["0"][0],
|
|
1.5 / 1.5, places=5)
|
|
self.assertAlmostEqual(extracted["magnitude_samples_unscaled"]["1"][0],
|
|
0.9 / 0.9, places=5)
|
|
self.assertEqual(extracted["per_lora_norm_sq"]["0"], 2.25)
|
|
self.assertEqual(extracted["target_key"], "layer.weight")
|
|
self.assertFalse(extracted["is_clip"])
|
|
self.assertEqual(extracted["strength_signs"]["0"], 1)
|
|
|
|
def test_reconstruct_rescales_by_new_strength(self):
|
|
"""Reconstruction rescales magnitude samples and display_l2 to new strength."""
|
|
cached_prefix = {
|
|
"pair_conflicts": {
|
|
"0,1": {"overlap": 50, "conflict": 10, "dot": 0.4,
|
|
"norm_a_sq": 2.25, "norm_b_sq": 0.81,
|
|
"weighted_total": 0.6, "weighted_conflict": 0.1,
|
|
"expected_conflict": 0.12, "excess_conflict": 0.0,
|
|
"subspace_overlap": 0.2, "subspace_weight": 1.35}
|
|
},
|
|
"per_lora_norm_sq": {"0": 2.25, "1": 0.81},
|
|
"magnitude_samples_unscaled": {"0": [1.0, 2.0], "1": [1.0]},
|
|
"ranks": {"0": 16, "1": 32},
|
|
"target_key": "layer.weight",
|
|
"is_clip": False,
|
|
"raw_n": 2,
|
|
"skip_count": 0,
|
|
"strength_signs": {"0": 1, "1": 1},
|
|
}
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 2.0, "clip_strength": None},
|
|
{"name": "b.safetensors", "strength": 0.5, "clip_strength": None},
|
|
]
|
|
result = lora_optimizer.LoRAOptimizer._reconstruct_from_analysis_cache(
|
|
"prefix_x", cached_prefix, active_loras)
|
|
self.assertIsNotNone(result)
|
|
prefix, partial_stats, pair_conflicts, mag_samples, target_info, skip_count, raw_n, norm_sq = result
|
|
|
|
self.assertEqual(prefix, "prefix_x")
|
|
# display_l2 = sqrt(norm_sq) * abs(strength)
|
|
self.assertAlmostEqual(partial_stats[0][2], math.sqrt(2.25) * 2.0, places=4)
|
|
self.assertAlmostEqual(partial_stats[1][2], math.sqrt(0.81) * 0.5, places=4)
|
|
# pair_conflicts keys are int tuples
|
|
self.assertIn((0, 1), pair_conflicts)
|
|
# magnitude_samples rescaled by abs(new_strength)
|
|
self.assertAlmostEqual(mag_samples[0][0].item(), 1.0 * 2.0, places=4)
|
|
self.assertAlmostEqual(mag_samples[1][0].item(), 1.0 * 0.5, places=4)
|
|
|
|
def test_reconstruct_returns_none_on_sign_flip(self):
|
|
"""Sign flip triggers None so caller falls back to full analysis."""
|
|
cached_prefix = {
|
|
"pair_conflicts": {},
|
|
"per_lora_norm_sq": {"0": 1.0},
|
|
"magnitude_samples_unscaled": {"0": [1.0]},
|
|
"ranks": {"0": 16},
|
|
"target_key": "layer.weight",
|
|
"is_clip": False,
|
|
"raw_n": 1,
|
|
"skip_count": 0,
|
|
"strength_signs": {"0": 1}, # was positive
|
|
}
|
|
active_loras = [{"name": "a.safetensors", "strength": -1.0, "clip_strength": None}]
|
|
result = lora_optimizer.LoRAOptimizer._reconstruct_from_analysis_cache(
|
|
"prefix_x", cached_prefix, active_loras)
|
|
self.assertIsNone(result)
|
|
|
|
def test_run_group_analysis_skips_analyze_on_cache_hit(self):
|
|
"""When cached_analysis covers a prefix, _analyze_target_group is not called."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
model = _make_model()
|
|
target_groups = optimizer._build_target_groups(
|
|
["prefix_a"], {"prefix_a": "layer.weight"}, {})
|
|
|
|
fake_cached = {
|
|
"prefix_a": {
|
|
"pair_conflicts": {
|
|
"0,1": {"overlap": 10, "conflict": 2, "dot": 0.1,
|
|
"norm_a_sq": 1.0, "norm_b_sq": 0.25,
|
|
"weighted_total": 0.3, "weighted_conflict": 0.05,
|
|
"expected_conflict": 0.1, "excess_conflict": 0.0,
|
|
"subspace_overlap": 0.1, "subspace_weight": 0.5}
|
|
},
|
|
"per_lora_norm_sq": {"0": 1.0, "1": 0.25},
|
|
"magnitude_samples_unscaled": {"0": [0.5], "1": [0.3]},
|
|
"ranks": {"0": 1, "1": 1},
|
|
"target_key": "layer.weight",
|
|
"is_clip": False,
|
|
"raw_n": 2,
|
|
"skip_count": 0,
|
|
"strength_signs": {"0": 1, "1": 1},
|
|
}
|
|
}
|
|
|
|
call_count = {"n": 0}
|
|
orig = optimizer._analyze_target_group
|
|
def counting_analyze(*args, **kwargs):
|
|
call_count["n"] += 1
|
|
return orig(*args, **kwargs)
|
|
|
|
with mock.patch.object(optimizer, "_analyze_target_group",
|
|
side_effect=counting_analyze):
|
|
device = torch.device("cpu")
|
|
result = optimizer._run_group_analysis(
|
|
target_groups, active_loras, model, None, device,
|
|
cached_analysis=fake_cached,
|
|
track_new_entries=True,
|
|
)
|
|
|
|
self.assertEqual(call_count["n"], 0)
|
|
self.assertEqual(result["new_analysis_entries"], {})
|
|
|
|
def test_run_group_analysis_populates_new_entries_on_miss(self):
|
|
"""Cache miss populates new_analysis_entries in the return dict."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
model = _make_model()
|
|
target_groups = optimizer._build_target_groups(
|
|
["prefix_a"], {"prefix_a": "layer.weight"}, {})
|
|
|
|
device = torch.device("cpu")
|
|
result = optimizer._run_group_analysis(
|
|
target_groups, active_loras, model, None, device,
|
|
cached_analysis={}, # empty cache = full miss
|
|
track_new_entries=True,
|
|
)
|
|
self.assertIn("prefix_a", result["new_analysis_entries"])
|
|
|
|
def test_run_group_analysis_calls_on_prefix_done_for_fresh_prefixes(self):
|
|
"""on_prefix_done fires once per freshly-computed prefix."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
model = _make_model()
|
|
target_groups = optimizer._build_target_groups(
|
|
["prefix_a"], {"prefix_a": "layer.weight"}, {})
|
|
|
|
calls = []
|
|
device = torch.device("cpu")
|
|
optimizer._run_group_analysis(
|
|
target_groups, active_loras, model, None, device,
|
|
cached_analysis={},
|
|
track_new_entries=True,
|
|
on_prefix_done=lambda prefix, entry: calls.append((prefix, entry)),
|
|
)
|
|
self.assertEqual(len(calls), 1)
|
|
prefix, entry = calls[0]
|
|
self.assertEqual(prefix, "prefix_a")
|
|
self.assertIn("ranks", entry)
|
|
|
|
def test_run_group_analysis_does_not_call_on_prefix_done_for_cache_hits(self):
|
|
"""on_prefix_done is NOT called for prefixes already in cached_analysis."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
model = _make_model()
|
|
target_groups = optimizer._build_target_groups(
|
|
["prefix_a"], {"prefix_a": "layer.weight"}, {})
|
|
|
|
fake_cached = {
|
|
"prefix_a": {
|
|
"pair_conflicts": {
|
|
"0,1": {"overlap": 10, "conflict": 2, "dot": 0.1,
|
|
"norm_a_sq": 1.0, "norm_b_sq": 0.25,
|
|
"weighted_total": 0.3, "weighted_conflict": 0.05,
|
|
"expected_conflict": 0.1, "excess_conflict": 0.0,
|
|
"subspace_overlap": 0.1, "subspace_weight": 0.5}
|
|
},
|
|
"per_lora_norm_sq": {"0": 1.0, "1": 0.25},
|
|
"magnitude_samples_unscaled": {"0": [0.5], "1": [0.3]},
|
|
"ranks": {"0": 1, "1": 1},
|
|
"target_key": "layer.weight",
|
|
"is_clip": False,
|
|
"raw_n": 2,
|
|
"skip_count": 0,
|
|
"strength_signs": {"0": 1, "1": 1},
|
|
}
|
|
}
|
|
|
|
calls = []
|
|
device = torch.device("cpu")
|
|
optimizer._run_group_analysis(
|
|
target_groups, active_loras, model, None, device,
|
|
cached_analysis=fake_cached,
|
|
track_new_entries=True,
|
|
on_prefix_done=lambda prefix, entry: calls.append(prefix),
|
|
)
|
|
self.assertEqual(calls, [])
|
|
|
|
def test_dataset_entry_includes_raw_analysis(self):
|
|
"""Dataset entries include raw_analysis field when new_analysis_entries provided."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.folder_paths.get_user_directory",
|
|
return_value=tmpdir):
|
|
tuner = lora_optimizer.LoRAAutoTuner()
|
|
tuner_data = {
|
|
"top_n": [],
|
|
"analysis_summary": {
|
|
"n_loras": 2, "prefix_count": 1,
|
|
"avg_conflict_ratio": 0.3, "avg_excess_conflict": 0.1,
|
|
"avg_subspace_overlap": 0.2, "avg_cosine_sim": 0.5,
|
|
"magnitude_ratio": 1.2, "decision_smoothing": 0.25,
|
|
},
|
|
"architecture_preset": "auto",
|
|
"lora_hash": "abc",
|
|
}
|
|
prefix_stats = {
|
|
"prefix_a": {
|
|
"n_loras": 2, "conflict_ratio": 0.3,
|
|
"excess_conflict": 0.1, "avg_cos_sim": 0.5,
|
|
"magnitude_ratio": 1.2, "avg_subspace_overlap": 0.2,
|
|
"magnitude_samples": [],
|
|
"per_lora_norm_sq": {0: 1.0, 1: 0.5},
|
|
"pairwise_dots": {},
|
|
"raw_n_loras": 2,
|
|
}
|
|
}
|
|
new_analysis_entries = {
|
|
"prefix_a": {
|
|
"pair_conflicts": {"0,1": {"overlap": 10}},
|
|
"per_lora_norm_sq": {"0": 1.0, "1": 0.5},
|
|
"magnitude_samples_unscaled": {"0": [0.5], "1": [0.3]},
|
|
"ranks": {"0": 1, "1": 1},
|
|
"target_key": "layer.weight",
|
|
"is_clip": False,
|
|
"raw_n": 2,
|
|
"skip_count": 0,
|
|
"strength_signs": {"0": 1, "1": 1},
|
|
}
|
|
}
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 1.0},
|
|
{"name": "b.safetensors", "strength": 0.5},
|
|
]
|
|
tuner._save_tuner_dataset_entry(
|
|
tuner_data, active_loras, prefix_stats, "wan_video",
|
|
names_only_hash="testhash",
|
|
new_analysis_entries=new_analysis_entries)
|
|
|
|
dataset_path = os.path.join(
|
|
tmpdir, "lora_optimizer_reports", "autotuner_dataset.jsonl")
|
|
with open(dataset_path) as f:
|
|
entry = json.loads(f.readline())
|
|
self.assertIn("raw_analysis", entry)
|
|
self.assertEqual(entry["raw_analysis"]["names_only_hash"], "testhash")
|
|
self.assertIn("prefix_a", entry["raw_analysis"]["per_prefix"])
|
|
self.assertIn("pair_conflicts",
|
|
entry["raw_analysis"]["per_prefix"]["prefix_a"])
|
|
|
|
|
|
class TestAnalysisPartialLifecycle(unittest.TestCase):
|
|
"""Integration tests for partial checkpoint create/resume/delete lifecycle."""
|
|
|
|
def test_partial_file_created_after_each_prefix(self):
|
|
"""Callback writes partial file; file exists after analysis."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
tuner = lora_optimizer.LoRAAutoTuner()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
names_only_hash, _ = tuner._compute_names_only_hash(active_loras)
|
|
source_loras = [{"name": l["name"]} for l in active_loras]
|
|
|
|
partial_accumulated = {}
|
|
written = []
|
|
|
|
def on_prefix_done(prefix, entry):
|
|
partial_accumulated[prefix] = entry
|
|
tuner._analysis_partial_save(names_only_hash, partial_accumulated, source_loras)
|
|
written.append(prefix)
|
|
|
|
model = _make_model()
|
|
target_groups = tuner._build_target_groups(
|
|
["prefix_a"], {"prefix_a": "layer.weight"}, {})
|
|
device = torch.device("cpu")
|
|
tuner._run_group_analysis(
|
|
target_groups, active_loras, model, None, device,
|
|
cached_analysis={},
|
|
track_new_entries=True,
|
|
on_prefix_done=on_prefix_done,
|
|
)
|
|
|
|
self.assertIn("prefix_a", written)
|
|
partial_path = tuner._analysis_partial_path(names_only_hash)
|
|
self.assertTrue(os.path.exists(partial_path))
|
|
|
|
def test_partial_file_deleted_after_full_cache_saved(self):
|
|
"""_analysis_partial_delete removes the file after full cache is written."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
tuner = lora_optimizer.LoRAAutoTuner()
|
|
hash_val = "testhash99"
|
|
tuner._analysis_partial_save(hash_val, {"prefix_a": {}}, [])
|
|
partial_path = tuner._analysis_partial_path(hash_val)
|
|
self.assertTrue(os.path.exists(partial_path))
|
|
|
|
tuner._analysis_cache_save(hash_val, {"prefix_a": {}}, [])
|
|
tuner._analysis_partial_delete(hash_val)
|
|
|
|
self.assertFalse(os.path.exists(partial_path))
|
|
full_path = tuner._analysis_cache_path(hash_val)
|
|
self.assertTrue(os.path.exists(full_path))
|
|
|
|
def test_partial_file_loaded_when_no_full_cache(self):
|
|
"""If full cache is missing but .partial exists, it is loaded."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
tuner = lora_optimizer.LoRAAutoTuner()
|
|
hash_val = "resumehash1"
|
|
per_prefix = {"prefix_b": {"ranks": {"0": 16}, "is_clip": False}}
|
|
tuner._analysis_partial_save(hash_val, per_prefix, [])
|
|
|
|
self.assertIsNone(tuner._analysis_cache_load(hash_val))
|
|
loaded = tuner._analysis_partial_load(hash_val)
|
|
self.assertIsNotNone(loaded)
|
|
self.assertIn("prefix_b", loaded)
|
|
|
|
def test_partial_promoted_to_full_cache_when_all_prefixes_already_done(self):
|
|
"""If partial has all prefixes and analysis adds none, full cache is still written."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
tuner = lora_optimizer.LoRAAutoTuner()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
names_only_hash, _ = tuner._compute_names_only_hash(active_loras)
|
|
source_loras = [{"name": l["name"]} for l in active_loras]
|
|
|
|
# Simulate: partial file has all prefixes (crash after last prefix, before full save)
|
|
model = _make_model()
|
|
target_groups = tuner._build_target_groups(
|
|
["prefix_a"], {"prefix_a": "layer.weight"}, {})
|
|
device = torch.device("cpu")
|
|
|
|
# Ensure clean state: remove any stale full cache from prior runs
|
|
full_path = tuner._analysis_cache_path(names_only_hash)
|
|
try:
|
|
os.unlink(full_path)
|
|
except OSError:
|
|
pass
|
|
tuner._analysis_partial_delete(names_only_hash)
|
|
|
|
# First: compute full analysis and manually save only to partial
|
|
result = tuner._run_group_analysis(
|
|
target_groups, active_loras, model, None, device,
|
|
cached_analysis={},
|
|
track_new_entries=True,
|
|
)
|
|
partial_entries = result["new_analysis_entries"]
|
|
tuner._analysis_partial_save(names_only_hash, partial_entries, source_loras)
|
|
|
|
# Verify: full cache does NOT exist yet (only partial was saved)
|
|
self.assertFalse(os.path.exists(full_path))
|
|
|
|
# Now simulate resume: load from partial (all prefixes cached), run analysis
|
|
cached_analysis = tuner._analysis_partial_load(names_only_hash)
|
|
self.assertIsNotNone(cached_analysis)
|
|
|
|
using_partial = True
|
|
partial_accumulated = dict(cached_analysis)
|
|
|
|
def on_prefix_done(prefix, entry):
|
|
partial_accumulated[prefix] = entry
|
|
tuner._analysis_partial_save(names_only_hash, partial_accumulated, source_loras)
|
|
|
|
result2 = tuner._run_group_analysis(
|
|
target_groups, active_loras, model, None, device,
|
|
cached_analysis=cached_analysis,
|
|
track_new_entries=True,
|
|
on_prefix_done=on_prefix_done,
|
|
)
|
|
new_entries = result2["new_analysis_entries"]
|
|
# All prefixes were cached, so new_entries should be empty
|
|
self.assertEqual(new_entries, {})
|
|
|
|
# Apply the fix: promote partial to full cache when using_partial=True
|
|
if new_entries or using_partial:
|
|
merged = dict(cached_analysis or {})
|
|
merged.update(new_entries)
|
|
tuner._analysis_cache_save(names_only_hash, merged, source_loras)
|
|
tuner._analysis_partial_delete(names_only_hash)
|
|
|
|
# Verify: full cache now exists, partial is gone
|
|
self.assertTrue(os.path.exists(full_path))
|
|
partial_path = tuner._analysis_partial_path(names_only_hash)
|
|
self.assertFalse(os.path.exists(partial_path))
|
|
|
|
|
|
class TestLoraCacheIO(unittest.TestCase):
|
|
|
|
def test_lora_identity_hash_returns_16char_hex(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Create a fake lora file so os.stat works
|
|
lora_path = os.path.join(tmpdir, "test.safetensors")
|
|
with open(lora_path, "w") as f:
|
|
f.write("x")
|
|
with mock.patch("lora_optimizer.folder_paths.get_full_path",
|
|
return_value=lora_path):
|
|
h = lora_optimizer.LoRAAutoTuner._lora_identity_hash(
|
|
{"name": "test.safetensors", "strength": 1.0})
|
|
self.assertIsInstance(h, str)
|
|
self.assertEqual(len(h), 16)
|
|
self.assertTrue(all(c in "0123456789abcdef" for c in h))
|
|
|
|
def test_lora_identity_hash_differs_for_different_files(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path_a = os.path.join(tmpdir, "a.safetensors")
|
|
path_b = os.path.join(tmpdir, "b.safetensors")
|
|
with open(path_a, "w") as f: f.write("x")
|
|
with open(path_b, "w") as f: f.write("y")
|
|
def fake_get_full_path(folder, name):
|
|
return path_a if name == "a.safetensors" else path_b
|
|
with mock.patch("lora_optimizer.folder_paths.get_full_path",
|
|
side_effect=fake_get_full_path):
|
|
ha = lora_optimizer.LoRAAutoTuner._lora_identity_hash(
|
|
{"name": "a.safetensors", "strength": 1.0})
|
|
hb = lora_optimizer.LoRAAutoTuner._lora_identity_hash(
|
|
{"name": "b.safetensors", "strength": 1.0})
|
|
self.assertNotEqual(ha, hb)
|
|
|
|
def test_lora_identity_hash_ignores_strength(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
lora_path = os.path.join(tmpdir, "test.safetensors")
|
|
with open(lora_path, "w") as f: f.write("x")
|
|
with mock.patch("lora_optimizer.folder_paths.get_full_path",
|
|
return_value=lora_path):
|
|
h1 = lora_optimizer.LoRAAutoTuner._lora_identity_hash(
|
|
{"name": "test.safetensors", "strength": 0.5})
|
|
h2 = lora_optimizer.LoRAAutoTuner._lora_identity_hash(
|
|
{"name": "test.safetensors", "strength": 1.0})
|
|
self.assertEqual(h1, h2)
|
|
|
|
def test_lora_cache_path_uses_lora_suffix(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
path = lora_optimizer.LoRAAutoTuner._lora_cache_path("abc123")
|
|
self.assertTrue(path.endswith("abc123.lora.json"))
|
|
|
|
def test_lora_cache_load_missing_returns_none(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
self.assertIsNone(
|
|
lora_optimizer.LoRAAutoTuner._lora_cache_load("nonexistent"))
|
|
|
|
def test_lora_cache_load_stale_algo_version_returns_none(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
path = os.path.join(tmpdir, "stale.lora.json")
|
|
with open(path, "w") as f:
|
|
json.dump({"algo_version": "0.0.0", "per_prefix": {}}, f)
|
|
self.assertIsNone(
|
|
lora_optimizer.LoRAAutoTuner._lora_cache_load("stale"))
|
|
|
|
def test_lora_cache_roundtrip(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
per_prefix = {
|
|
"prefix_a": {
|
|
"norm_sq": 1.5, "rank": 16,
|
|
"magnitude_samples_unscaled": [0.1, 0.2],
|
|
"strength_sign": 1,
|
|
"target_key": "layer.weight",
|
|
"is_clip": False, "skip_count": 0, "raw_n": 1,
|
|
}
|
|
}
|
|
lora_optimizer.LoRAAutoTuner._lora_cache_save("abc123", per_prefix)
|
|
loaded = lora_optimizer.LoRAAutoTuner._lora_cache_load("abc123")
|
|
self.assertIsNotNone(loaded)
|
|
self.assertIn("prefix_a", loaded)
|
|
self.assertEqual(loaded["prefix_a"]["norm_sq"], 1.5)
|
|
|
|
def test_lora_cache_save_is_atomic(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
lora_optimizer.LoRAAutoTuner._lora_cache_save("abc123", {})
|
|
path = lora_optimizer.LoRAAutoTuner._lora_cache_path("abc123")
|
|
self.assertTrue(os.path.exists(path))
|
|
self.assertFalse(os.path.exists(path + ".tmp"))
|
|
|
|
|
|
class TestPairCacheIO(unittest.TestCase):
|
|
|
|
def test_pair_cache_path_sorts_hashes(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
path_ab = lora_optimizer.LoRAAutoTuner._pair_cache_path("aaa", "bbb")
|
|
path_ba = lora_optimizer.LoRAAutoTuner._pair_cache_path("bbb", "aaa")
|
|
self.assertEqual(path_ab, path_ba)
|
|
self.assertIn("aaa_bbb", path_ab)
|
|
|
|
def test_pair_cache_load_missing_returns_none(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
self.assertIsNone(
|
|
lora_optimizer.LoRAAutoTuner._pair_cache_load("missing1", "missing2"))
|
|
|
|
def test_pair_cache_load_stale_algo_version_returns_none(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
path = os.path.join(tmpdir, "stale1_stale2.pair.json")
|
|
with open(path, "w") as f:
|
|
json.dump({"algo_version": "0.0.0", "per_prefix": {}}, f)
|
|
self.assertIsNone(
|
|
lora_optimizer.LoRAAutoTuner._pair_cache_load("stale1", "stale2"))
|
|
|
|
def test_pair_cache_roundtrip(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
per_prefix = {
|
|
"prefix_a": {
|
|
"overlap": 100, "conflict": 30, "dot": 0.5,
|
|
"norm_a_sq": 1.0, "norm_b_sq": 0.5,
|
|
"weighted_total": 0.8, "weighted_conflict": 0.2,
|
|
"expected_conflict": 0.15, "excess_conflict": 0.05,
|
|
"subspace_overlap": 0.3, "subspace_weight": 1.0,
|
|
}
|
|
}
|
|
lora_optimizer.LoRAAutoTuner._pair_cache_save("aaa", "bbb", per_prefix)
|
|
loaded = lora_optimizer.LoRAAutoTuner._pair_cache_load("aaa", "bbb")
|
|
self.assertIsNotNone(loaded)
|
|
self.assertEqual(loaded["prefix_a"]["overlap"], 100)
|
|
|
|
def test_pair_cache_load_commutative(self):
|
|
"""_pair_cache_load("a","b") and ("b","a") return the same data."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
per_prefix = {"prefix_a": {"overlap": 50}}
|
|
lora_optimizer.LoRAAutoTuner._pair_cache_save("aaa", "bbb", per_prefix)
|
|
loaded_ab = lora_optimizer.LoRAAutoTuner._pair_cache_load("aaa", "bbb")
|
|
loaded_ba = lora_optimizer.LoRAAutoTuner._pair_cache_load("bbb", "aaa")
|
|
self.assertEqual(loaded_ab, loaded_ba)
|
|
|
|
def test_pair_cache_save_is_atomic(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
lora_optimizer.LoRAAutoTuner._pair_cache_save("aaa", "bbb", {})
|
|
path = lora_optimizer.LoRAAutoTuner._pair_cache_path("aaa", "bbb")
|
|
self.assertTrue(os.path.exists(path))
|
|
self.assertFalse(os.path.exists(path + ".tmp"))
|
|
|
|
|
|
class TestCacheExtraction(unittest.TestCase):
|
|
|
|
def _make_result(self):
|
|
"""Minimal 8-tuple from _analyze_target_group."""
|
|
partial_stats = [(0, 16, 1.5, 2.25), (1, 32, 0.9, 0.81)]
|
|
pair_conflicts = {
|
|
(0, 1): {
|
|
"overlap": 50, "conflict": 10, "dot": 0.4,
|
|
"norm_a_sq": 2.25, "norm_b_sq": 0.81,
|
|
"weighted_total": 0.6, "weighted_conflict": 0.1,
|
|
"expected_conflict": 0.12, "excess_conflict": 0.0,
|
|
"subspace_overlap": 0.2, "subspace_weight": 0.5,
|
|
}
|
|
}
|
|
magnitude_samples = [
|
|
torch.tensor([0.5, 1.0]), # lora 0, already scaled by strength
|
|
torch.tensor([0.3, 0.6]), # lora 1
|
|
]
|
|
per_lora_norm_sq = {0: 2.25, 1: 0.81}
|
|
return (
|
|
"prefix_a", partial_stats, pair_conflicts, magnitude_samples,
|
|
("layer.weight", False), 0, 2, per_lora_norm_sq
|
|
)
|
|
|
|
def test_extract_for_lora_cache_participating(self):
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 1.5},
|
|
{"name": "b.safetensors", "strength": 0.9},
|
|
]
|
|
result = self._make_result()
|
|
entry = lora_optimizer.LoRAOptimizer._extract_for_lora_cache(result, 0, active_loras)
|
|
self.assertIsNotNone(entry)
|
|
self.assertEqual(entry["norm_sq"], 2.25)
|
|
self.assertEqual(entry["rank"], 16)
|
|
self.assertEqual(entry["strength_sign"], 1)
|
|
self.assertFalse(entry["is_clip"])
|
|
self.assertEqual(entry["target_key"], "layer.weight")
|
|
# magnitude unscaled: tensor / abs(strength=1.5)
|
|
self.assertAlmostEqual(entry["magnitude_samples_unscaled"][0], 0.5 / 1.5, places=5)
|
|
|
|
def test_extract_for_lora_cache_non_participating_returns_none(self):
|
|
"""LoRA index not in per_lora_norm_sq → non-participating → return None."""
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 1.0},
|
|
{"name": "b.safetensors", "strength": 1.0},
|
|
{"name": "c.safetensors", "strength": 1.0}, # index 2: not in result
|
|
]
|
|
result = self._make_result()
|
|
entry = lora_optimizer.LoRAOptimizer._extract_for_lora_cache(result, 2, active_loras)
|
|
self.assertIsNone(entry)
|
|
|
|
def test_extract_for_pair_cache_norm_order_by_hash(self):
|
|
"""norm_a_sq in entry corresponds to the LoRA with the smaller hash."""
|
|
result = self._make_result()
|
|
# hash_i > hash_j → swap norm_a/norm_b
|
|
entry = lora_optimizer.LoRAOptimizer._extract_for_pair_cache(
|
|
result, i=0, j=1, hash_i="zzz", hash_j="aaa")
|
|
# In result, norm_a_sq=2.25 (lora 0), norm_b_sq=0.81 (lora 1)
|
|
# hash_j="aaa" < hash_i="zzz", so "aaa" (lora 1) should be norm_a
|
|
self.assertAlmostEqual(entry["norm_a_sq"], 0.81)
|
|
self.assertAlmostEqual(entry["norm_b_sq"], 2.25)
|
|
|
|
def test_extract_for_pair_cache_no_swap_when_hash_i_smaller(self):
|
|
result = self._make_result()
|
|
entry = lora_optimizer.LoRAOptimizer._extract_for_pair_cache(
|
|
result, i=0, j=1, hash_i="aaa", hash_j="zzz")
|
|
# hash_i="aaa" < hash_j="zzz" → no swap
|
|
self.assertAlmostEqual(entry["norm_a_sq"], 2.25)
|
|
self.assertAlmostEqual(entry["norm_b_sq"], 0.81)
|
|
|
|
def test_extract_for_pair_cache_non_participating_returns_none(self):
|
|
"""Pair not in pair_conflicts (one LoRA doesn't participate) → None."""
|
|
result = self._make_result()
|
|
# Pair (0,2) not in result
|
|
entry = lora_optimizer.LoRAOptimizer._extract_for_pair_cache(
|
|
result, i=0, j=2, hash_i="aaa", hash_j="zzz")
|
|
self.assertIsNone(entry)
|
|
|
|
|
|
class TestPairLoraReconstruction(unittest.TestCase):
|
|
|
|
def _make_lora_entries(self):
|
|
return {
|
|
0: {
|
|
"norm_sq": 2.25, "rank": 16,
|
|
"magnitude_samples_unscaled": [0.5, 1.0],
|
|
"strength_sign": 1,
|
|
"target_key": "layer.weight",
|
|
"is_clip": False, "skip_count": 0, "raw_n": 2,
|
|
},
|
|
1: {
|
|
"norm_sq": 0.81, "rank": 32,
|
|
"magnitude_samples_unscaled": [0.3, 0.6],
|
|
"strength_sign": 1,
|
|
"target_key": "layer.weight",
|
|
"is_clip": False, "skip_count": 0, "raw_n": 2,
|
|
},
|
|
}
|
|
|
|
def _make_pair_entries(self, hash_0="aaa", hash_1="bbb"):
|
|
# hash_0 < hash_1, so norm_a_sq = lora 0
|
|
return {
|
|
(0, 1): {
|
|
"overlap": 50, "conflict": 10, "dot": 0.4,
|
|
"norm_a_sq": 2.25, "norm_b_sq": 0.81,
|
|
"weighted_total": 0.6, "weighted_conflict": 0.1,
|
|
"expected_conflict": 0.12, "excess_conflict": 0.0,
|
|
"subspace_overlap": 0.2, "subspace_weight": 0.5,
|
|
}
|
|
}
|
|
|
|
def test_reconstruction_returns_8tuple(self):
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 1.5},
|
|
{"name": "b.safetensors", "strength": 0.9},
|
|
]
|
|
lora_hashes = {0: "aaa", 1: "bbb"}
|
|
result = lora_optimizer.LoRAOptimizer._reconstruct_from_pair_lora_cache(
|
|
"prefix_a",
|
|
lora_entries=self._make_lora_entries(),
|
|
pair_entries=self._make_pair_entries(),
|
|
active_loras=active_loras,
|
|
lora_hashes=lora_hashes,
|
|
)
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(len(result), 8)
|
|
prefix, partial_stats, pair_conflicts, magnitude_samples, target_info, skip_count, raw_n, per_lora_norm_sq = result
|
|
self.assertEqual(prefix, "prefix_a")
|
|
self.assertEqual(len(partial_stats), 2)
|
|
self.assertIn((0, 1), pair_conflicts)
|
|
self.assertAlmostEqual(per_lora_norm_sq[0], 2.25)
|
|
|
|
def test_reconstruction_rescales_magnitude_by_current_strength(self):
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 2.0}, # different from cache
|
|
{"name": "b.safetensors", "strength": 1.0},
|
|
]
|
|
lora_hashes = {0: "aaa", 1: "bbb"}
|
|
result = lora_optimizer.LoRAOptimizer._reconstruct_from_pair_lora_cache(
|
|
"prefix_a",
|
|
lora_entries=self._make_lora_entries(),
|
|
pair_entries=self._make_pair_entries(),
|
|
active_loras=active_loras,
|
|
lora_hashes=lora_hashes,
|
|
)
|
|
_, _, _, magnitude_samples, _, _, _, _ = result
|
|
# unscaled[0] = 0.5, rescaled by strength=2.0 → 1.0
|
|
self.assertAlmostEqual(magnitude_samples[0][0].item(), 1.0, places=5)
|
|
|
|
def test_reconstruction_returns_none_on_sign_flip(self):
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": -1.0}, # sign flipped vs cached +1
|
|
{"name": "b.safetensors", "strength": 1.0},
|
|
]
|
|
lora_hashes = {0: "aaa", 1: "bbb"}
|
|
result = lora_optimizer.LoRAOptimizer._reconstruct_from_pair_lora_cache(
|
|
"prefix_a",
|
|
lora_entries=self._make_lora_entries(),
|
|
pair_entries=self._make_pair_entries(),
|
|
active_loras=active_loras,
|
|
lora_hashes=lora_hashes,
|
|
)
|
|
self.assertIsNone(result)
|
|
|
|
def test_reconstruction_returns_none_on_clip_strength_sign_flip(self):
|
|
"""Sign flip detected via clip_strength for CLIP prefixes."""
|
|
lora_entries = {
|
|
0: {
|
|
"norm_sq": 2.25, "rank": 16,
|
|
"magnitude_samples_unscaled": [0.5],
|
|
"strength_sign": -1, # cached with negative clip_strength
|
|
"target_key": "lora_te_text_model_encoder_layers_0_weight",
|
|
"is_clip": True, "skip_count": 0, "raw_n": 1,
|
|
},
|
|
1: {**self._make_lora_entries()[1],
|
|
"is_clip": True},
|
|
}
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 1.0, "clip_strength": 0.5}, # sign flip vs -1
|
|
{"name": "b.safetensors", "strength": 1.0},
|
|
]
|
|
lora_hashes = {0: "aaa", 1: "bbb"}
|
|
result = lora_optimizer.LoRAOptimizer._reconstruct_from_pair_lora_cache(
|
|
"prefix_clip",
|
|
lora_entries=lora_entries,
|
|
pair_entries=self._make_pair_entries(),
|
|
active_loras=active_loras,
|
|
lora_hashes=lora_hashes,
|
|
)
|
|
self.assertIsNone(result)
|
|
|
|
def test_reconstruction_swaps_norm_when_hash_ordering_differs(self):
|
|
"""When lora 0's hash > lora 1's hash, norm_a_sq in pair file is lora 1's."""
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 1.0},
|
|
{"name": "b.safetensors", "strength": 1.0},
|
|
]
|
|
lora_hashes = {0: "zzz", 1: "aaa"} # lora 1 has smaller hash
|
|
# In the pair file, norm_a_sq = lora with smaller hash = lora 1 = 0.81
|
|
pair_entries = {
|
|
(0, 1): {
|
|
"overlap": 50, "conflict": 10, "dot": 0.4,
|
|
"norm_a_sq": 0.81, # lora 1 (smaller hash "aaa")
|
|
"norm_b_sq": 2.25, # lora 0 (larger hash "zzz")
|
|
"weighted_total": 0.6, "weighted_conflict": 0.1,
|
|
"expected_conflict": 0.12, "excess_conflict": 0.0,
|
|
"subspace_overlap": 0.2, "subspace_weight": 0.5,
|
|
}
|
|
}
|
|
result = lora_optimizer.LoRAOptimizer._reconstruct_from_pair_lora_cache(
|
|
"prefix_a",
|
|
lora_entries=self._make_lora_entries(),
|
|
pair_entries=pair_entries,
|
|
active_loras=active_loras,
|
|
lora_hashes=lora_hashes,
|
|
)
|
|
self.assertIsNotNone(result)
|
|
_, _, pair_conflicts, _, _, _, _, _ = result
|
|
# After reconstruction, norm_a_sq should be for lora 0 = 2.25
|
|
self.assertAlmostEqual(pair_conflicts[(0, 1)]["norm_a_sq"], 2.25)
|
|
self.assertAlmostEqual(pair_conflicts[(0, 1)]["norm_b_sq"], 0.81)
|
|
|
|
def test_reconstruction_handles_non_participating_loras(self):
|
|
"""LoRA with None entry (non-participating) is excluded from partial_stats."""
|
|
active_loras = [
|
|
{"name": "a.safetensors", "strength": 1.0},
|
|
{"name": "b.safetensors", "strength": 1.0},
|
|
{"name": "c.safetensors", "strength": 1.0}, # non-participating
|
|
]
|
|
lora_hashes = {0: "aaa", 1: "bbb", 2: "ccc"}
|
|
lora_entries = {**self._make_lora_entries(), 2: None}
|
|
pair_entries = self._make_pair_entries() # only (0,1)
|
|
result = lora_optimizer.LoRAOptimizer._reconstruct_from_pair_lora_cache(
|
|
"prefix_a",
|
|
lora_entries=lora_entries,
|
|
pair_entries=pair_entries,
|
|
active_loras=active_loras,
|
|
lora_hashes=lora_hashes,
|
|
)
|
|
self.assertIsNotNone(result)
|
|
_, partial_stats, pair_conflicts, _, _, _, _, _ = result
|
|
lora_indices_in_stats = [s[0] for s in partial_stats]
|
|
self.assertNotIn(2, lora_indices_in_stats)
|
|
self.assertNotIn((0, 2), pair_conflicts)
|
|
self.assertNotIn((1, 2), pair_conflicts)
|
|
|
|
|
|
class TestPairLoraCacheWiring(unittest.TestCase):
|
|
|
|
def _make_lora_caches(self, prefix="prefix_a"):
|
|
"""Minimal lora_caches dict covering prefix_a for 2 loras."""
|
|
entry = {
|
|
"norm_sq": 1.0, "rank": 1,
|
|
"magnitude_samples_unscaled": [0.5],
|
|
"strength_sign": 1,
|
|
"target_key": "layer.weight",
|
|
"is_clip": False, "skip_count": 0, "raw_n": 2,
|
|
}
|
|
return {
|
|
0: {prefix: entry},
|
|
1: {prefix: {**entry, "norm_sq": 0.25}},
|
|
}
|
|
|
|
def _make_pair_caches(self, prefix="prefix_a"):
|
|
return {
|
|
(0, 1): {
|
|
prefix: {
|
|
"overlap": 10, "conflict": 2, "dot": 0.1,
|
|
"norm_a_sq": 1.0, "norm_b_sq": 0.25,
|
|
"weighted_total": 0.3, "weighted_conflict": 0.05,
|
|
"expected_conflict": 0.1, "excess_conflict": 0.0,
|
|
"subspace_overlap": 0.1, "subspace_weight": 0.5,
|
|
}
|
|
}
|
|
}
|
|
|
|
def test_run_group_analysis_uses_pair_lora_cache_on_full_hit(self):
|
|
"""When pair+lora caches cover all prefixes, _analyze_target_group is not called."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
model = _make_model()
|
|
target_groups = optimizer._build_target_groups(
|
|
["prefix_a"], {"prefix_a": "layer.weight"}, {})
|
|
|
|
call_count = {"n": 0}
|
|
orig = optimizer._analyze_target_group
|
|
def counting(*args, **kwargs):
|
|
call_count["n"] += 1
|
|
return orig(*args, **kwargs)
|
|
|
|
lora_hashes = {0: "aaa", 1: "bbb"}
|
|
with mock.patch.object(optimizer, "_analyze_target_group", side_effect=counting):
|
|
result = optimizer._run_group_analysis(
|
|
target_groups, active_loras, model, None, torch.device("cpu"),
|
|
lora_caches=self._make_lora_caches(),
|
|
pair_caches=self._make_pair_caches(),
|
|
lora_hashes=lora_hashes,
|
|
track_new_entries=True,
|
|
)
|
|
self.assertEqual(call_count["n"], 0)
|
|
self.assertEqual(result["new_lora_entries"][0], {})
|
|
self.assertEqual(result["new_pair_entries"][(0, 1)], {})
|
|
|
|
def test_run_group_analysis_populates_new_entries_on_pair_lora_miss(self):
|
|
"""Cache miss populates new_lora_entries and new_pair_entries."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
model = _make_model()
|
|
target_groups = optimizer._build_target_groups(
|
|
["prefix_a"], {"prefix_a": "layer.weight"}, {})
|
|
lora_hashes = {0: "aaa", 1: "bbb"}
|
|
result = optimizer._run_group_analysis(
|
|
target_groups, active_loras, model, None, torch.device("cpu"),
|
|
lora_caches={0: {}, 1: {}}, # prefix_a missing → full miss
|
|
pair_caches={(0, 1): {}},
|
|
lora_hashes=lora_hashes,
|
|
track_new_entries=True,
|
|
)
|
|
self.assertIn("prefix_a", result["new_lora_entries"][0])
|
|
self.assertIn("prefix_a", result["new_pair_entries"][(0, 1)])
|
|
|
|
|
|
class TestPairLoraCacheAutoTune(unittest.TestCase):
|
|
"""Verify that auto_tune loads, uses, and saves pair/lora caches."""
|
|
|
|
def test_lora_and_pair_cache_files_created_after_analysis(self):
|
|
"""After _run_group_analysis with misses, lora and pair files are saved."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
tuner = lora_optimizer.LoRAAutoTuner()
|
|
active_loras = [
|
|
_make_lora_entry({"prefix_a": 1.0}, strength=1.0, name="a.safetensors"),
|
|
_make_lora_entry({"prefix_a": 0.5}, strength=1.0, name="b.safetensors"),
|
|
]
|
|
|
|
lora_hashes = {}
|
|
with mock.patch("lora_optimizer.folder_paths.get_full_path",
|
|
return_value=None):
|
|
for i, lora in enumerate(active_loras):
|
|
lora_hashes[i] = tuner._lora_identity_hash(lora)
|
|
|
|
new_lora_entries = {0: {"prefix_a": {"norm_sq": 1.0}},
|
|
1: {"prefix_a": {"norm_sq": 0.5}}}
|
|
new_pair_entries = {(0, 1): {"prefix_a": {"overlap": 10}}}
|
|
|
|
# Save as if auto_tune just completed
|
|
for i, h in lora_hashes.items():
|
|
tuner._lora_cache_save(h, new_lora_entries[i])
|
|
tuner._pair_cache_save(
|
|
lora_hashes[0], lora_hashes[1], new_pair_entries[(0, 1)])
|
|
|
|
# Verify files exist
|
|
path_0 = tuner._lora_cache_path(lora_hashes[0])
|
|
path_1 = tuner._lora_cache_path(lora_hashes[1])
|
|
path_pair = tuner._pair_cache_path(lora_hashes[0], lora_hashes[1])
|
|
self.assertTrue(os.path.exists(path_0))
|
|
self.assertTrue(os.path.exists(path_1))
|
|
self.assertTrue(os.path.exists(path_pair))
|
|
|
|
def test_lora_cache_merged_with_existing_on_save(self):
|
|
"""New per-prefix entries are merged into the existing lora cache file."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir):
|
|
tuner = lora_optimizer.LoRAAutoTuner()
|
|
# Pre-existing cache with prefix_a
|
|
tuner._lora_cache_save("hash_x", {"prefix_a": {"norm_sq": 1.0}})
|
|
# New analysis added prefix_b
|
|
existing = tuner._lora_cache_load("hash_x") or {}
|
|
existing["prefix_b"] = {"norm_sq": 2.0}
|
|
tuner._lora_cache_save("hash_x", existing)
|
|
# Both prefixes should be in the file
|
|
loaded = tuner._lora_cache_load("hash_x")
|
|
self.assertIn("prefix_a", loaded)
|
|
self.assertIn("prefix_b", loaded)
|
|
|
|
def test_score_merge_result_lora_adapter_on_device(self):
|
|
"""_score_merge_result should handle LoRAAdapter patches correctly."""
|
|
LoRAAdapter = lora_optimizer.LoRAAdapter
|
|
up = torch.randn(8, 4)
|
|
down = torch.randn(4, 16)
|
|
alpha = 4.0
|
|
adapter = LoRAAdapter(
|
|
loaded_keys=set(),
|
|
weights=(up, down, alpha, None, None, None),
|
|
)
|
|
patches = {("key1",): adapter}
|
|
result = lora_optimizer._score_merge_result(patches, {}, compute_svd=False)
|
|
self.assertIn("norm_mean", result)
|
|
self.assertGreater(result["norm_mean"], 0)
|
|
self.assertIn("composite_score", result)
|
|
|
|
def test_sample_pair_metrics_downsamples_large_vectors(self):
|
|
"""Pair metrics should work correctly with large vectors that trigger downsampling."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
a = torch.randn(200000)
|
|
b = torch.randn(200000)
|
|
result = optimizer._sample_pair_metrics(a, b)
|
|
self.assertIn("overlap", result)
|
|
self.assertIn("conflict", result)
|
|
self.assertIn("dot", result)
|
|
self.assertGreater(result["overlap"], 0)
|
|
result2 = optimizer._sample_pair_metrics(a, b)
|
|
self.assertEqual(result["overlap"], result2["overlap"])
|
|
self.assertEqual(result["conflict"], result2["conflict"])
|
|
self.assertAlmostEqual(result["dot"], result2["dot"], places=4)
|
|
|
|
|
|
def test_sl_patch_cache_populates_and_hits(self):
|
|
"""Single-LoRA patch cache should store results and reuse them on matching auto_strength."""
|
|
optimizer = lora_optimizer.LoRAOptimizer()
|
|
# Simulate a result tuple as returned by _merge_one_group
|
|
fake_patch = ("diff", (torch.randn(4, 4),))
|
|
fake_result = ("weight.key", False, fake_patch, "weighted_sum", "lora_unet_block", 0.0, 1, False, 1.0, 0.9)
|
|
|
|
cache = {}
|
|
prefix = "lora_unet_block"
|
|
auto_strength = "enabled"
|
|
|
|
# First access: miss → populate
|
|
key = (prefix, auto_strength)
|
|
self.assertNotIn(key, cache)
|
|
cache[key] = fake_result
|
|
self.assertIn(key, cache)
|
|
self.assertIs(cache[key], fake_result)
|
|
|
|
# Second access with same auto_strength: hit
|
|
self.assertIs(cache.get(key), fake_result)
|
|
|
|
# Different auto_strength: miss
|
|
key2 = (prefix, "disabled")
|
|
self.assertIsNone(cache.get(key2))
|
|
|
|
# Different prefix, same auto_strength: miss
|
|
key3 = ("lora_unet_other", auto_strength)
|
|
self.assertIsNone(cache.get(key3))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|