520 lines
19 KiB
Python
520 lines
19 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_torch = types.ModuleType("safetensors.torch")
|
|
safetensors_torch.save_file = lambda state_dict, path: None
|
|
safetensors.torch = safetensors_torch
|
|
sys.modules["safetensors"] = safetensors
|
|
sys.modules["safetensors.torch"] = safetensors_torch
|
|
|
|
|
|
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,
|
|
"activation_ratio": 1.5,
|
|
"per_lora_norm_sq": {0: 1.0, 1: 1.0},
|
|
"per_lora_activation_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,
|
|
"activation_ratio": 1.5,
|
|
"per_lora_norm_sq": {0: 1.0, 1: 1.0},
|
|
"per_lora_activation_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_activation_importance_uses_calibration_diag(self):
|
|
diff = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
|
|
calibration_entry = {"input_diag": [4.0, 0.0]}
|
|
|
|
importance = self.optimizer._compute_activation_importance(diff, calibration_entry)
|
|
|
|
self.assertAlmostEqual(importance, 4.0)
|
|
|
|
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_score_merge_result_uses_activation_importance(self):
|
|
patch_a = ("diff", (torch.tensor([[1.0, 0.0]], dtype=torch.float32),))
|
|
patch_b = ("diff", (torch.tensor([[0.0, 1.0]], dtype=torch.float32),))
|
|
|
|
metrics = lora_optimizer._score_merge_result(
|
|
{"layer.weight": patch_a, "layer2.weight": patch_b},
|
|
{},
|
|
compute_svd=False,
|
|
calibration_data={
|
|
"targets": {
|
|
"layer.weight": {"input_diag": [10.0, 0.0]},
|
|
"layer2.weight": {"input_diag": [0.0, 1.0]},
|
|
}
|
|
},
|
|
)
|
|
|
|
self.assertIn("importance_cv", metrics)
|
|
self.assertGreater(metrics["importance_mean"], 0.0)
|
|
|
|
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: captured.update({"state_dict": state_dict, "path": path})):
|
|
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()
|
|
calibration_saver = lora_optimizer.SaveCalibrationData()
|
|
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")
|
|
|
|
with self.assertRaises(ValueError):
|
|
calibration_saver.save_calibration_data({"targets": {}}, "../escape")
|
|
|
|
def test_save_and_load_calibration_data(self):
|
|
saver = lora_optimizer.SaveCalibrationData()
|
|
loader = lora_optimizer.LoadCalibrationData()
|
|
payload = {"targets": {"layer.weight": {"input_diag": [1.0, 2.0]}}}
|
|
|
|
path, = saver.save_calibration_data(payload, "tests/calibration_test")
|
|
self.assertTrue(path.endswith(".json"))
|
|
|
|
with mock.patch.object(
|
|
lora_optimizer.folder_paths,
|
|
"get_full_path_or_raise",
|
|
return_value=path,
|
|
):
|
|
loaded, = loader.load_calibration_data("ignored.json")
|
|
|
|
self.assertEqual(loaded, payload)
|
|
|
|
def test_bridge_passthrough_returns_ui_payload(self):
|
|
tuner_data = {
|
|
"top_n": [{
|
|
"config": {
|
|
"optimization_mode": "global",
|
|
"merge_mode": "ties",
|
|
"merge_quality": "standard",
|
|
"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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|