EthanfelandClaude Sonnet 4.6
aae488bc00
Adds module-level _extract_lora_svd() that SVD-decomposes a weight delta into (lora_up, lora_down, alpha), supporting fixed-rank and auto-rank (energy threshold) modes with a singular value noise floor. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1164 lines
50 KiB
Python
1164 lines
50 KiB
Python
import importlib.util
|
|
import json
|
|
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
|
|
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_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("record_dataset"), "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")
|
|
|
|
|
|
@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
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|