nodes/ethanfel-ComfyUI-LoRA-Optimizer
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
9ce78a8aa051a66bce87c0dacf142b50bc9cd092
ethanfel-ComfyUI-LoRA-Optim…/tests/test_inline_optimizer.py
T
EthanfelandClaude Fable 5 9ce78a8aa0 feat(inline): run the AutoTuner (search/memory/community) on captured chains 
execute_inline's autotuner-mode settings previously set autotuner_fallback
and merged with optimizer defaults ("not supported inline"). Now that
captured items have a stable content-based identity (prior commit),
delegate mode=="autotuner" to a LoRAAutoTuner exactly like
LoRAOptimizerSimple — but on the inline-prepared STRIPPED model/clip and
the VIRTUAL stack (never the raw model / original chain), with
normalize_keys forced "disabled" (captured keys are already canonical).
auto_tune runs candidates through optimize_merge, which already handles
_precomputed_diffs, and applies the winner off the passed stack (no file
reload). The 6-tuple return is reconciled to the inline 5-tuple with the
chain fingerprint prepended. Advanced-mode and no-settings paths are
unchanged; the fallback message now only fires for genuinely unknown modes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 17:34:58 +02:00

1760 lines
84 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Tests for the inline chain-filter optimizer node."""
import struct
import types
import unittest
import uuid
from unittest import mock
import torch
# Reuse the stub installer / module instance from the main test module.
from tests.test_lora_optimizer import lora_optimizer
_ALPHA_DEFAULT = object()
def _adapter(rank=4, out_dim=8, in_dim=8, up=None, down=None,
alpha=_ALPHA_DEFAULT):
"""Minimal LoRAAdapter-like payload the engine can expand. Explicit
up/down matrices make the expanded diff deterministic (alpha == rank, so
the diff is exactly up @ down). Pass alpha=None to mimic PEFT/diffusers/
AI-Toolkit LoRAs whose files carry no .alpha key — comfy stores such
adapters with alpha=None (scale 1.0)."""
if up is None:
up = torch.randn(out_dim, rank)
if down is None:
down = torch.randn(rank, in_dim)
if alpha is _ALPHA_DEFAULT:
alpha = float(rank) # alpha == rank -> scale 1.0, diff == up @ down
return lora_optimizer.LoRAAdapter(
loaded_keys=set(),
weights=(up, down, alpha, None, None, None),
)
def _entry(strength, payload, strength_model=1.0, offset=None, function=None):
"""A ModelPatcher patch-list entry, shaped like model_patcher.py:807."""
return (strength, payload, strength_model, offset, function)
class TestPatchClassification(unittest.TestCase):
def test_adapter_is_capturable(self):
e = _entry(0.8, _adapter())
self.assertTrue(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
def test_diff_tuple_is_capturable(self):
e = _entry(1.0, ("diff", (torch.randn(4, 4),)))
self.assertTrue(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
def test_set_tuple_passes_through(self):
e = _entry(1.0, ("set", (torch.randn(4, 4),)))
self.assertFalse(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
def test_nonunit_strength_model_passes_through(self):
e = _entry(1.0, _adapter(), strength_model=0.5)
self.assertFalse(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
def test_function_entry_passes_through(self):
e = _entry(1.0, _adapter(), function=lambda w: w)
self.assertFalse(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
def test_unknown_object_passes_through(self):
e = _entry(1.0, object())
self.assertFalse(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
def test_non_5_tuple_entry_passes_through(self):
# older/nonstandard third-party nodes may store short entries — the
# classifier must not crash on unpack, just pass them through
e = (1.0, _adapter(), 1.0)
self.assertFalse(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
def test_padded_diff_passes_through(self):
# comfy pads the BASE weight at apply time for this shape — we cannot
# faithfully expand it to a bare diff tensor
e = _entry(1.0, ("diff", (torch.randn(4, 4), {"pad_weight": True})))
self.assertFalse(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
def test_malformed_diff_passes_through(self):
e = _entry(1.0, ("diff", None))
self.assertFalse(lora_optimizer._LoRAMergeBase._is_capturable_entry(e))
class TestNoneAlphaExpansion(unittest.TestCase):
"""PEFT/diffusers/AI-Toolkit LoRAs have no .alpha key, so a captured comfy
LoRAAdapter carries alpha=None. _expand_patch_to_diff must treat that as
scale 1.0 (comfy semantics: scale = alpha/rank if alpha is not None else
1.0) instead of crashing on None / rank."""
def test_none_alpha_adapter_expands_at_scale_one(self):
up = torch.randn(8, 4)
down = torch.randn(4, 8)
adapter = _adapter(up=up, down=down, alpha=None)
got = lora_optimizer._LoRAMergeBase._expand_patch_to_diff(adapter)
self.assertTrue(torch.allclose(got, up @ down, atol=1e-6))
def test_numeric_alpha_still_rescales(self):
# regression net: a real alpha still divides by rank
up = torch.randn(8, 4)
down = torch.randn(4, 8)
adapter = _adapter(up=up, down=down, alpha=2.0) # rank 4 -> scale 0.5
got = lora_optimizer._LoRAMergeBase._expand_patch_to_diff(adapter)
self.assertTrue(torch.allclose(got, (up @ down) * 0.5, atol=1e-6))
def _chain_patches(*loras):
"""Simulate a loader chain: each lora is {key: payload} applied at a
strength. Distinct float objects per call, entries appended in order —
exactly what ModelPatcher.add_patches does."""
patches = {}
seen_ids = set()
for strength_value, lora in loras:
# struct round-trip mints a FRESH float object even for values whose
# literals CPython folds/interns (e.g. 1.0 used twice in one test).
# float(x) would return the SAME object for an exact float input.
s = struct.unpack("d", struct.pack("d", strength_value))[0]
assert id(s) not in seen_ids, "float object reused across simulated calls"
seen_ids.add(id(s))
for key, payload in lora.items():
patches.setdefault(key, []).append(_entry(s, payload))
return patches
class TestChainGroupReconstruction(unittest.TestCase):
def _reconstruct(self, patches):
return lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(patches)
def test_single_lora(self):
patches = _chain_patches((0.8, {"a": _adapter(), "b": _adapter()}))
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 1)
self.assertAlmostEqual(groups[0]["strength"], 0.8)
self.assertEqual(set(groups[0]["entries"]), {"a", "b"})
def test_two_loras_chain_order(self):
patches = _chain_patches(
(0.8, {"a": _adapter(), "b": _adapter()}),
(0.5, {"a": _adapter(), "c": _adapter()}),
)
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 2)
self.assertAlmostEqual(groups[0]["strength"], 0.8)
self.assertAlmostEqual(groups[1]["strength"], 0.5)
self.assertEqual(set(groups[1]["entries"]), {"a", "c"})
def test_subset_alignment(self):
# A patches attn only; B patches attn+mlp. Naive index-as-identity
# would misattribute B's mlp entry (position 0 there) to A.
patches = _chain_patches(
(0.8, {"attn": _adapter()}),
(0.5, {"attn": _adapter(), "mlp": _adapter()}),
)
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 2)
self.assertEqual(set(groups[0]["entries"]), {"attn"})
self.assertEqual(set(groups[1]["entries"]), {"attn", "mlp"})
def test_same_strength_value_distinct_objects(self):
patches = _chain_patches(
(1.0, {"a": _adapter()}),
(1.0, {"a": _adapter(), "b": _adapter()}),
)
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 2)
def test_shared_float_object_fallback(self):
# Pathological: two calls share ONE float object (interning). The
# id-group then holds two entries on key "a" — must split by per-key
# order instead of returning a corrupt group.
s = 1.0
patches = {}
for lora in ({"a": _adapter()}, {"a": _adapter(), "b": _adapter()}):
for key, payload in lora.items():
patches.setdefault(key, []).append(_entry(s, payload))
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 2)
for g in groups:
keys = [k for k in g["entries"]]
self.assertEqual(len(keys), len(set(keys)))
def test_offset_entries_get_tuple_keys(self):
off = (0, 0, 4)
patches = {"qkv": [_entry(0.7, _adapter(), offset=off)]}
groups = self._reconstruct(patches)
self.assertEqual(list(groups[0]["entries"]), [("qkv", off)])
def test_noncapturable_entries_ignored(self):
patches = _chain_patches((0.8, {"a": _adapter()}))
patches["a"].append(_entry(1.0, ("set", (torch.zeros(2, 2),))))
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 1)
self.assertEqual(set(groups[0]["entries"]), {"a"})
def test_empty_patches(self):
self.assertEqual(self._reconstruct({}), [])
def test_non_5_tuple_entry_ignored_without_raising(self):
patches = _chain_patches((0.8, {"a": _adapter()}))
patches["a"].append((1.0, _adapter(), 1.0)) # nonstandard 3-tuple
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 1)
self.assertEqual(set(groups[0]["entries"]), {"a"})
def test_interleaved_collision_no_fragmentation(self):
# Distinct-strength loader X on {a} first, then two interned-strength
# loaders B and C both patching {a, b}. The collision sub-gid must be
# the per-target-key collision ORDINAL — splitting on the absolute
# per-key position fragments this into 4 groups [a],[a,b],[a],[b]
# because X shifts the positions on "a" but not on "b".
d = struct.unpack("d", struct.pack("d", 0.6))[0]
s = 1.0
pX = _adapter()
pB_a, pB_b = _adapter(), _adapter()
pC_a, pC_b = _adapter(), _adapter()
patches = {
"a": [_entry(d, pX), _entry(s, pB_a), _entry(s, pC_a)],
"b": [_entry(s, pB_b), _entry(s, pC_b)],
}
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 3)
self.assertEqual(set(groups[0]["entries"]), {"a"})
self.assertEqual(set(groups[1]["entries"]), {"a", "b"})
self.assertEqual(set(groups[2]["entries"]), {"a", "b"})
# ordinal-k entries belong to the k-th colliding call on EVERY key
self.assertIs(groups[1]["entries"]["a"], pB_a)
self.assertIs(groups[1]["entries"]["b"], pB_b)
self.assertIs(groups[2]["entries"]["a"], pC_a)
self.assertIs(groups[2]["entries"]["b"], pC_b)
def test_collision_subgroup_reordered_by_precedence(self):
# Chain: two interned-strength calls A, B on "a", then distinct call D
# on {z, a}. Dict iteration starts at "z", so D's group is CREATED
# first; both collision groups (base + ordinal sub-group) must be
# sorted forward past D via the shared-key "a" positions — this
# exercises the insertion sort on a collision-created sub-group.
s = 1.0
d = struct.unpack("d", struct.pack("d", 0.5))[0]
pA, pB, pD_a, pD_z = _adapter(), _adapter(), _adapter(), _adapter()
patches = {
"z": [_entry(d, pD_z)],
"a": [_entry(s, pA), _entry(s, pB), _entry(d, pD_a)],
}
groups = self._reconstruct(patches)
self.assertEqual(len(groups), 3)
self.assertIs(groups[0]["entries"]["a"], pA)
self.assertIs(groups[1]["entries"]["a"], pB)
self.assertEqual(set(groups[2]["entries"]), {"z", "a"})
class _FakePatcher:
"""Minimal ModelPatcher stand-in: ordered patch lists + clone()."""
def __init__(self, patches=None):
self.patches = patches if patches is not None else {}
self.patches_uuid = object()
def clone(self):
return _FakePatcher({k: v[:] for k, v in self.patches.items()})
class TestStripCaptured(unittest.TestCase):
def test_strips_only_captured_entries(self):
keep = _entry(1.0, ("set", (torch.zeros(2, 2),)))
patches = _chain_patches((0.8, {"a": _adapter(), "b": _adapter()}))
patches["a"].append(keep)
patcher = _FakePatcher(patches)
groups = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(patcher.patches)
clone = patcher.clone()
lora_optimizer._LoRAMergeBase._strip_captured_entries(clone, groups)
self.assertEqual(clone.patches, {"a": [keep]})
self.assertEqual(len(patcher.patches["a"]), 2) # original untouched
def test_uuid_regenerated(self):
patcher = _FakePatcher(_chain_patches((0.8, {"a": _adapter()})))
groups = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(patcher.patches)
clone = patcher.clone()
before = clone.patches_uuid
lora_optimizer._LoRAMergeBase._strip_captured_entries(clone, groups)
self.assertNotEqual(clone.patches_uuid, before)
def _slot(enabled=True, strength=1.0, model_strength=1.0, clip_strength=1.0,
conflict_mode="all", key_filter="all", preserve=False):
return dict(enabled=enabled, strength=strength, model_strength=model_strength,
clip_strength=clip_strength, conflict_mode=conflict_mode,
key_filter=key_filter, preserve=preserve)
class TestChainStackBuild(unittest.TestCase):
def _build(self, model_groups, clip_groups, slots, visibility="simple"):
return lora_optimizer.LoRAOptimizerInline._chain_groups_to_stack(
model_groups, clip_groups, slots, visibility)
def test_basic_item_schema(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"a": _adapter()})))
stack = self._build(mg, [], [_slot()])
item = stack[0]
self.assertTrue(item["_precomputed_diffs"])
self.assertAlmostEqual(item["strength"], 0.8) # loader strength kept
self.assertIsNone(item["clip_strength"])
self.assertEqual(item["conflict_mode"], "all")
self.assertIn("a", item["lora"])
def test_simple_mode_multiplier(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"a": _adapter()})))
stack = self._build(mg, [], [_slot(strength=0.5)])
self.assertAlmostEqual(stack[0]["strength"], 0.4) # 0.8 loader × 0.5 slot
def test_advanced_mode_split_multipliers(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"a": _adapter()})))
cg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.6, {"te.a": _adapter()})))
stack = self._build(mg, cg, [_slot(model_strength=0.5, clip_strength=2.0)],
visibility="advanced")
self.assertAlmostEqual(stack[0]["strength"], 0.4) # 0.8 × 0.5
self.assertAlmostEqual(stack[0]["clip_strength"], 1.2) # 0.6 × 2.0
self.assertIn("te.a", stack[0]["lora"]) # clip keys merged in
def test_disabled_slot_excluded(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"b": _adapter()})))
stack = self._build(mg, [], [_slot(enabled=False), _slot()])
self.assertEqual(len(stack), 1)
self.assertAlmostEqual(stack[0]["strength"], 0.5)
def test_missing_slots_get_defaults(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"b": _adapter()})))
stack = self._build(mg, [], [_slot(preserve=True)]) # only 1 slot for 2 loras
self.assertEqual(len(stack), 2)
self.assertTrue(stack[0]["preserve"])
self.assertFalse(stack[1]["preserve"])
def test_leftover_clip_groups_become_clip_only_items(self):
cg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.6, {"te.a": _adapter()})))
stack = self._build([], cg, [])
self.assertEqual(len(stack), 1)
self.assertIn("te.a", stack[0]["lora"])
def test_leftover_strength_is_clip_product_nonzero(self):
# the engine drops strength == 0 items (active_loras filter in
# optimize_merge) — clip-only items must ride on the clip product
cg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.6, {"te.a": _adapter()})))
stack = self._build([], cg, [])
self.assertNotEqual(stack[0]["strength"], 0.0)
self.assertAlmostEqual(stack[0]["strength"], 0.6)
self.assertAlmostEqual(stack[0]["clip_strength"], 0.6)
def test_zeroed_model_branch_keeps_live_clip_branch(self):
# stock-UI case: LoraLoader with strength_model=0, strength_clip=1.0
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"a": _adapter()})))
cg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.6, {"te.a": _adapter()})))
stack = self._build(mg, cg, [_slot(model_strength=0.0, clip_strength=1.0)],
visibility="advanced")
self.assertEqual(len(stack), 1)
self.assertNotIn("a", stack[0]["lora"]) # unet keys dropped
self.assertIn("te.a", stack[0]["lora"])
self.assertAlmostEqual(stack[0]["strength"], 0.6) # clip product
self.assertAlmostEqual(stack[0]["clip_strength"], 0.6)
def test_both_branches_zeroed_item_excluded(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"a": _adapter()})))
cg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.6, {"te.a": _adapter()})))
stack = self._build(mg, cg, [_slot(model_strength=0.0, clip_strength=0.0)],
visibility="advanced")
self.assertEqual(stack, [])
def test_leftover_after_model_groups_uses_own_slot_index(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"a": _adapter()})))
cg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.6, {"te.a": _adapter()}), (0.4, {"te.b": _adapter()})))
stack = self._build(mg, cg, [_slot(), _slot()])
self.assertEqual(len(stack), 2)
self.assertEqual(stack[1]["name"], "chain lora #2 (clip-only)")
self.assertIn("te.b", stack[1]["lora"])
self.assertAlmostEqual(stack[1]["strength"], 0.4)
# slot index j (not 0) gates the leftover — disabling slot 2 drops it
stack = self._build(mg, cg, [_slot(), _slot(enabled=False)])
self.assertEqual(len(stack), 1)
self.assertNotIn("(clip-only)", stack[0]["name"])
def test_simple_visibility_scales_paired_clip(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"a": _adapter()})))
cg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.6, {"te.a": _adapter()})))
stack = self._build(mg, cg, [_slot(strength=0.5)]) # simple visibility
self.assertAlmostEqual(stack[0]["strength"], 0.4) # 0.8 × 0.5
self.assertAlmostEqual(stack[0]["clip_strength"], 0.3) # 0.6 × 0.5
def test_extra_slots_ignored(self):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"a": _adapter()})))
stack = self._build(mg, [], [_slot(), _slot(strength=9.0), _slot(enabled=False)])
self.assertEqual(len(stack), 1)
self.assertAlmostEqual(stack[0]["strength"], 0.8)
class _FakeCLIP:
def __init__(self, patcher):
self.patcher = patcher
self.cond_stage_model = types.SimpleNamespace()
def clone(self):
return _FakeCLIP(self.patcher.clone())
class TestInlineExecute(unittest.TestCase):
def _node(self):
return lora_optimizer.LoRAOptimizerInline()
def _model(self, patches):
m = _FakePatcher(patches)
m.model = types.SimpleNamespace()
return m
@staticmethod
def _out(result):
return result["result"] if isinstance(result, dict) else result
def _capture_merge(self, node, seen):
"""Mock optimize_merge, recording every kwarg + the stack/model."""
def fake_merge(m, stack, output_strength, **kw):
seen["model"] = m
seen["stack"] = stack
seen["output_strength"] = output_strength
seen.update(kw)
return (m, kw.get("clip"), "engine report", None, None)
node.optimize_merge = fake_merge
def test_no_lora_patches_passthrough(self):
model = self._model({})
node = self._node()
result = node.execute_inline(model, output_strength=1.0)
out = self._out(result)
self.assertIs(out[0], model) # unchanged object back
self.assertIn("No LoRA patches", out[2])
self.assertIn("AFTER your Load LoRA nodes", out[2])
self.assertIsNone(out[3])
self.assertIsNone(out[4])
def test_non_mergeable_only_patches_reported_honestly(self):
# Patches exist but none are capturable (OFT/"set"/function entries):
# "place this node AFTER your loaders" would be WRONG advice.
patches = {"a": [_entry(1.0, ("set", (torch.zeros(2, 2),))),
_entry(1.0, _adapter(), function=lambda w: w)]}
model = self._model(patches)
node = self._node()
result = node.execute_inline(model, output_strength=1.0)
out = self._out(result)
self.assertIs(out[0], model)
self.assertIn("2 non-mergeable patch entries", out[2])
self.assertIn("passed through", out[2])
self.assertNotIn("AFTER your Load LoRA nodes", out[2])
def test_merge_called_with_virtual_stack_and_stripped_model(self):
model = self._model(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"a": _adapter()})))
node = self._node()
seen = {}
self._capture_merge(node, seen)
node.execute_inline(model, output_strength=1.0)
self.assertEqual(len(seen["stack"]), 2)
self.assertTrue(all(i["_precomputed_diffs"] for i in seen["stack"]))
self.assertEqual(seen["model"].patches, {}) # stripped clone
self.assertIsNot(seen["model"], model)
self.assertEqual(seen["normalize_keys"], "disabled")
def test_report_prepends_fingerprints(self):
model = self._model(_chain_patches((0.8, {"a": _adapter(), "b": _adapter()})))
node = self._node()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "engine report", None, None)
result = node.execute_inline(model, output_strength=1.0)
report = self._out(result)[2]
self.assertIn("#1", report)
self.assertIn("2 keys", report)
self.assertIn("0.80", report)
self.assertIn("rank 4", report)
self.assertIn("engine report", report)
def test_disabled_slot_stripped_but_not_merged(self):
model = self._model(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"b": _adapter()})))
node = self._node()
seen = {}
self._capture_merge(node, seen)
node.execute_inline(model, output_strength=1.0,
chain_options={"visibility": "simple",
"slots": [_slot(enabled=False), _slot()]})
self.assertEqual(len(seen["stack"]), 1)
self.assertAlmostEqual(seen["stack"][0]["strength"], 0.5)
self.assertEqual(seen["model"].patches, {}) # disabled LoRA stripped too
def test_unconnected_chain_options_merges_all_with_defaults(self):
# chain_options unconnected (None): every captured LoRA merges with
# default options, and the report shows NO slots-vs-LoRAs note
# (merge-all-defaults is not a mismatch).
model = self._model(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"b": _adapter()})))
node = self._node()
seen = {}
self._capture_merge(node, seen)
report = self._out(node.execute_inline(model, output_strength=1.0))[2]
self.assertEqual(len(seen["stack"]), 2) # both merged
self.assertAlmostEqual(seen["stack"][0]["strength"], 0.8) # default mult
self.assertAlmostEqual(seen["stack"][1]["strength"], 0.5)
self.assertFalse(any(i.get("preserve") for i in seen["stack"]))
self.assertNotIn("option slot", report) # no slots-vs-LoRAs
self.assertEqual(seen["normalize_keys"], "disabled")
def test_fingerprint_warns_on_group_count_mismatch(self):
# 2 model groups vs 1 clip group (both non-empty) -> order-attribution
# warning naming the usual suspects
model = self._model(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"b": _adapter()})))
clip = _FakeCLIP(_FakePatcher(_chain_patches((0.6, {"te.a": _adapter()}))))
node = self._node()
seen = {}
self._capture_merge(node, seen)
result = node.execute_inline(model, output_strength=1.0, clip=clip)
report = self._out(result)[2]
self.assertIn("2 model-side vs 1 clip-side", report)
self.assertIn("LoraLoaderModelOnly", report)
# the clip clone passed to the engine is stripped, original untouched
self.assertIsNot(seen["clip"], clip)
self.assertEqual(seen["clip"].patcher.patches, {})
self.assertEqual(len(clip.patcher.patches["te.a"]), 1)
def test_fingerprint_counts_passthrough_entries(self):
patches = _chain_patches((0.8, {"a": _adapter()}))
patches["a"].append(_entry(1.0, ("set", (torch.zeros(2, 2),))))
model = self._model(patches)
node = self._node()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "r", None, None)
result = node.execute_inline(model, output_strength=1.0)
report = self._out(result)[2]
self.assertIn("1 non-LoRA", report)
self.assertIn("passed through untouched", report)
def test_no_passthrough_note_when_all_captured(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "r", None, None)
report = self._out(node.execute_inline(model, output_strength=1.0))[2]
self.assertNotIn("non-LoRA", report)
def test_fingerprint_notes_fewer_slots_than_loras(self):
model = self._model(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"b": _adapter()})))
node = self._node()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "r", None, None)
report = self._out(node.execute_inline(
model, output_strength=1.0,
chain_options={"visibility": "simple", "slots": [_slot()]}))[2]
self.assertIn("2 LoRAs detected but only 1 option slot", report)
self.assertIn("default options", report)
def test_fingerprint_notes_more_slots_than_loras(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "r", None, None)
report = self._out(node.execute_inline(
model, output_strength=1.0,
chain_options={"visibility": "simple",
"slots": [_slot(), _slot(), _slot()]}))[2]
self.assertIn("3 option slots but only 1 LoRA", report)
self.assertIn("extra slots ignored", report)
def test_advanced_settings_delegated_normalize_keys_pinned(self):
settings = {
"mode": "advanced",
"auto_strength": "disabled", "auto_strength_floor": 0.5,
"optimization_mode": "global",
"sparsification": "dare", "sparsification_density": 0.5,
"dare_dampening": 0.1,
"merge_refinement": "refine", "strategy_set": "basic",
"normalize_keys": "enabled", # must NOT reach the engine
"architecture_preset": "dit", "decision_smoothing": 0.1,
"smooth_slerp_gate": True,
"star_eta": 0.9, "tame_layers": 0.1, "tame_threshold": 0.2,
"cache_patches": "disabled", "patch_compression": "disabled",
"svd_device": "cpu", "free_vram_between_passes": "enabled",
"vram_budget": 0.25,
}
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
seen = {}
self._capture_merge(node, seen)
result = node.execute_inline(model, output_strength=1.0, settings=settings)
report = self._out(result)[2]
self.assertNotIn("not supported inline", report)
self.assertEqual(seen["normalize_keys"], "disabled") # pinned
self.assertEqual(seen["optimization_mode"], "global")
self.assertEqual(seen["sparsification"], "dare")
self.assertAlmostEqual(seen["sparsification_density"], 0.5)
self.assertEqual(seen["cache_patches"], "disabled")
self.assertEqual(seen["svd_device"], "cpu")
self.assertAlmostEqual(seen["star_eta"], 0.9)
# explicit preset from the settings node -> no arch-unknown warning
self.assertNotIn("architecture: unknown", report)
def test_fingerprint_warns_arch_unknown_without_preset(self):
# Virtual items skip _normalize_stack's arch detection, so auto
# detection can never resolve — the report must say so.
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "r", None, None)
report = self._out(node.execute_inline(model, output_strength=1.0))[2]
self.assertIn("architecture: unknown (inline capture)", report)
self.assertIn("architecture_preset", report)
def test_passthrough_count_includes_clip_side(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
clip_patches = _chain_patches((0.6, {"te.a": _adapter()}))
clip_patches["te.a"].append(_entry(1.0, ("set", (torch.zeros(2, 2),))))
clip = _FakeCLIP(_FakePatcher(clip_patches))
node = self._node()
node.optimize_merge = lambda m, s, o, **kw: (m, kw.get("clip"), "r", None, None)
report = self._out(node.execute_inline(model, output_strength=1.0,
clip=clip))[2]
self.assertIn("1 non-LoRA", report) # the clip-side "set" entry
self.assertIn("passed through untouched", report)
def test_lokr_payload_reports_dense_rank(self):
# LoKr weights[1] is the w2 Kronecker factor, NOT a rank — the
# fingerprint must not misreport its shape as "rank N".
lokr = lora_optimizer.LoKrAdapter(
loaded_keys=set(),
weights=(torch.randn(4, 4), torch.randn(2, 2), None, None,
None, None, None, None))
model = self._model(_chain_patches((0.8, {"a": lokr})))
node = self._node()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "r", None, None)
report = self._out(node.execute_inline(model, output_strength=1.0))[2]
self.assertIn("dense", report)
self.assertNotIn("rank 2", report)
def test_advanced_visibility_widgets_reach_stack(self):
model = self._model(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"b": _adapter()})))
clip = _FakeCLIP(_FakePatcher(_chain_patches(
(0.6, {"te.a": _adapter()}), (0.4, {"te.b": _adapter()}))))
node = self._node()
seen = {}
self._capture_merge(node, seen)
node.execute_inline(model, output_strength=1.0, clip=clip,
chain_options={"visibility": "advanced",
"slots": [_slot(),
_slot(model_strength=0.5,
clip_strength=2.0)]})
self.assertEqual(len(seen["stack"]), 2)
self.assertAlmostEqual(seen["stack"][0]["strength"], 0.8) # defaults
self.assertAlmostEqual(seen["stack"][0]["clip_strength"], 0.6)
self.assertAlmostEqual(seen["stack"][1]["strength"], 0.25) # 0.5 x 0.5
self.assertAlmostEqual(seen["stack"][1]["clip_strength"], 0.8) # 0.4 x 2.0
def test_salt_names_differ_across_source_patchers(self):
# Two DIFFERENT loader chains at identical strengths must yield
# different virtual item names, or optimize_merge's in-node cache
# would false-hit when the user swaps a LoRA file.
names_per_run = []
for _ in range(2):
model = self._model(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"b": _adapter()})))
model.patches_uuid = uuid.uuid4()
node = self._node()
seen = {}
self._capture_merge(node, seen)
node.execute_inline(model, output_strength=1.0)
names_per_run.append([i["name"] for i in seen["stack"]])
self.assertNotEqual(names_per_run[0], names_per_run[1])
for names in names_per_run:
self.assertTrue(names[0].startswith("chain lora #1"))
self.assertTrue(names[1].startswith("chain lora #2"))
def test_salt_names_stable_for_same_patcher(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
model.patches_uuid = uuid.uuid4()
node = self._node()
names_per_run = []
def fake_merge(m, stack, o, **kw):
names_per_run.append([i["name"] for i in stack])
return (m, None, "r", None, None)
node.optimize_merge = fake_merge
node.execute_inline(model, output_strength=1.0)
node.execute_inline(model, output_strength=1.0)
self.assertEqual(names_per_run[0], names_per_run[1])
def test_input_types_surface(self):
cls = lora_optimizer.LoRAOptimizerInline
it = cls.INPUT_TYPES()
req, opt = it["required"], it["optional"]
# widgets the user types live in "required"
for w in ("model", "output_strength", "clip_strength_multiplier"):
self.assertIn(w, req)
# per-LoRA option widgets moved to LoRAInlineChainOptions — the inline
# node no longer carries settings_visibility/lora_count or any slot
self.assertNotIn("settings_visibility", req)
self.assertNotIn("lora_count", req)
for base in ("enabled", "strength", "model_strength", "clip_strength",
"conflict_mode", "key_filter", "preserve"):
self.assertNotIn(f"{base}_1", req)
# genuine node-to-node wires stay optional
self.assertIn("clip", opt)
self.assertIn("settings", opt)
self.assertIn("chain_options", opt) # new side-node input
self.assertEqual(cls.FUNCTION, "execute_inline")
# inherited from LoRAOptimizer — lora_data keeps SaveMergedLoRA working
self.assertEqual(cls.RETURN_TYPES,
("MODEL", "CLIP", "STRING", "TUNER_DATA", "LORA_DATA"))
self.assertNotIn("RETURN_TYPES", vars(cls)) # not redeclared
class TestInlineIsChanged(unittest.TestCase):
"""LoRAOptimizerInline.IS_CHANGED must accept the node's ACTUAL (now much
smaller) inputs — the inherited LoRAOptimizer signature expects lora_stack
and would raise -> node re-merges on every queue press — and key on the
upstream patcher state + chain_options content + settings."""
def _model(self):
m = _FakePatcher(_chain_patches((0.8, {"a": _adapter()})))
m.patches_uuid = uuid.uuid4()
return m
def test_accepts_actual_inputs_and_is_stable(self):
cls = lora_optimizer.LoRAOptimizerInline
model = self._model()
first = cls.IS_CHANGED(model, 1.0) # must not raise
self.assertEqual(first, cls.IS_CHANGED(model, 1.0))
def test_changes_when_chain_options_content_changes(self):
cls = lora_optimizer.LoRAOptimizerInline
model = self._model()
co = {"visibility": "simple", "slots": [_slot()]}
first = cls.IS_CHANGED(model, 1.0, chain_options=co)
self.assertEqual(first, cls.IS_CHANGED(model, 1.0, chain_options=co))
changed = {"visibility": "simple", "slots": [_slot(enabled=False)]}
self.assertNotEqual(first,
cls.IS_CHANGED(model, 1.0, chain_options=changed))
def test_changes_when_upstream_chain_changes(self):
cls = lora_optimizer.LoRAOptimizerInline
model = self._model()
first = cls.IS_CHANGED(model, 1.0)
model.patches_uuid = uuid.uuid4() # loader chain re-executed
self.assertNotEqual(first, cls.IS_CHANGED(model, 1.0))
def test_changes_with_output_strength_and_clip_multiplier(self):
cls = lora_optimizer.LoRAOptimizerInline
model = self._model()
first = cls.IS_CHANGED(model, 1.0)
self.assertNotEqual(first, cls.IS_CHANGED(model, 0.5))
self.assertNotEqual(first, cls.IS_CHANGED(
model, 1.0, clip_strength_multiplier=0.5))
def test_changes_with_clip_and_settings(self):
cls = lora_optimizer.LoRAOptimizerInline
model = self._model()
first = cls.IS_CHANGED(model, 1.0)
clip = _FakeCLIP(_FakePatcher(_chain_patches((0.6, {"te.a": _adapter()}))))
clip.patcher.patches_uuid = uuid.uuid4()
with_clip = cls.IS_CHANGED(model, 1.0, clip=clip)
self.assertNotEqual(first, with_clip)
clip.patcher.patches_uuid = uuid.uuid4()
self.assertNotEqual(with_clip, cls.IS_CHANGED(model, 1.0, clip=clip))
self.assertNotEqual(first, cls.IS_CHANGED(
model, 1.0, settings={"mode": "advanced"}))
def _realistic_load_lora_for_models(calls):
"""A comfy.sd.load_lora_for_models stand-in that mimics the REAL failure
mode instead of blindly passing the model through: it resolves only
TRAINER-format keys ({x}.lora_up.weight / .lora_down.weight). Virtual
chain items carry model-target keys, so nothing matches and the models
come back unchanged — exactly what live comfy does, silently."""
def stub(model, clip, lora_dict, strength_model, strength_clip):
matched = [k for k in lora_dict
if isinstance(k, str) and (k.endswith(".lora_up.weight")
or k.endswith(".lora_down.weight"))]
calls.append({"matched": matched, "n_keys": len(lora_dict)})
return (model, clip) # nothing baked for unmatched keys
return stub
def _pipeline_model(patches, applied, dim=16):
"""Fake ModelPatcher threading the REAL optimize_merge pipeline: a .model
with a dim x dim base weight so _resolve_target_shape/_resolve_base_norm
work, faithful clone() (copies the CURRENT patch lists, like the real
ModelPatcher.clone — so a clone of the stripped clone stays stripped),
and an add_patches recorder capturing the merged patches the engine
re-applies (plus which patcher received them)."""
base = types.SimpleNamespace(
layer=types.SimpleNamespace(weight=torch.zeros(dim, dim)))
def _attach(p):
p.model = base
p.add_patches = lambda patches_, strength=1.0, strength_clip=None: (
applied.update(patches=dict(patches_), strength=strength,
patcher=p),
list(patches_.keys()))[1]
p.clone = lambda: _attach(
_FakePatcher({k: v[:] for k, v in p.patches.items()}))
return p
return _attach(_FakePatcher(patches))
class TestSingleLoraVirtualPath(unittest.TestCase):
"""A 1-LoRA chain (or N LoRAs with all but one disabled) must NOT take
optimize_merge's single-LoRA fast path: load_lora_for_models looks up
trainer-format keys, virtual items carry model-target keys, so it would
load NOTHING — and with the originals stripped, the output model would
silently lose the LoRA entirely."""
def _fake_model(self, patches, applied):
return _pipeline_model(patches, applied)
def test_single_virtual_lora_goes_through_pipeline(self):
key = "layer.weight"
applied = {}
model = self._fake_model(
_chain_patches((0.9, {key: _adapter(rank=4, out_dim=16, in_dim=16)})),
applied)
node = lora_optimizer.LoRAOptimizerInline()
node._get_model_keys = lambda m: {"alias_layer": key}
calls = []
with mock.patch.object(lora_optimizer.comfy.sd, "load_lora_for_models",
_realistic_load_lora_for_models(calls)):
result = node.execute_inline(model, output_strength=1.0)
out = result["result"] if isinstance(result, dict) else result
self.assertEqual(calls, []) # fast path NOT taken
self.assertTrue(applied.get("patches")) # pipeline re-applied patches
self.assertIn("chain lora #1", out[2])
def test_disabled_down_to_single_also_uses_pipeline(self):
key = "layer.weight"
applied = {}
model = self._fake_model(
_chain_patches((0.9, {key: _adapter(rank=4, out_dim=16, in_dim=16)}),
(0.5, {key: _adapter(rank=4, out_dim=16, in_dim=16)})),
applied)
node = lora_optimizer.LoRAOptimizerInline()
node._get_model_keys = lambda m: {"alias_layer": key}
calls = []
with mock.patch.object(lora_optimizer.comfy.sd, "load_lora_for_models",
_realistic_load_lora_for_models(calls)):
node.execute_inline(model, output_strength=1.0,
chain_options={"visibility": "simple",
"slots": [_slot(),
_slot(enabled=False)]})
self.assertEqual(calls, [])
self.assertTrue(applied.get("patches"))
def test_plain_single_lora_stack_still_takes_fast_path(self):
# Regression net for the guard: trainer-format single-LoRA stacks
# keep the fast path, and the realistic stub CAN match their keys.
model = _FakePatcher({})
model.model = types.SimpleNamespace(
layer=types.SimpleNamespace(weight=torch.zeros(16, 16)))
stack = [{"name": "A", "strength": 1.0,
"lora": {"alias_layer.lora_up.weight": torch.randn(16, 4),
"alias_layer.lora_down.weight": torch.randn(4, 16),
"alias_layer.alpha": torch.tensor(4.0)}}]
node = lora_optimizer.LoRAOptimizer()
node._get_model_keys = lambda m: {"alias_layer": "layer.weight"}
calls = []
with mock.patch.object(lora_optimizer.comfy.sd, "load_lora_for_models",
_realistic_load_lora_for_models(calls)):
node.optimize_merge(model, stack, 1.0)
self.assertEqual(len(calls), 1)
self.assertTrue(calls[0]["matched"])
class TestChainOptionsNode(unittest.TestCase):
"""The side node carries the per-LoRA widgets and emits a
LORA_CHAIN_OPTIONS payload the inline node consumes by chain order."""
def _node(self):
return lora_optimizer.LoRAInlineChainOptions()
@staticmethod
def _full_slot_kwargs():
"""One distinct value per widget across all 10 slots so both the
per-slot mapping and lora_count truncation are observable."""
cls = lora_optimizer.LoRAInlineChainOptions
kw = {}
for i in range(1, cls.MAX_LORAS + 1):
kw[f"enabled_{i}"] = (i % 2 == 0)
kw[f"strength_{i}"] = float(i)
kw[f"model_strength_{i}"] = float(i) + 0.1
kw[f"clip_strength_{i}"] = float(i) + 0.2
kw[f"conflict_mode_{i}"] = "all"
kw[f"key_filter_{i}"] = "all"
kw[f"preserve_{i}"] = (i % 2 == 1)
return kw
def test_returns_visibility_passthrough_and_slot_count(self):
node = self._node()
(out,) = node.build_options("advanced", 3, **self._full_slot_kwargs())
self.assertEqual(out["visibility"], "advanced")
self.assertEqual(len(out["slots"]), 3)
def test_slot_dicts_have_all_seven_keys_with_passed_values(self):
node = self._node()
(out,) = node.build_options("simple", 2, **self._full_slot_kwargs())
expected_keys = set(lora_optimizer.LoRAOptimizerInline._SLOT_DEFAULTS)
self.assertEqual(len(expected_keys), 7)
for i, slot in enumerate(out["slots"], start=1):
self.assertEqual(set(slot), expected_keys)
self.assertEqual(slot["strength"], float(i))
self.assertEqual(slot["model_strength"], float(i) + 0.1)
self.assertEqual(slot["clip_strength"], float(i) + 0.2)
self.assertEqual(slot["enabled"], (i % 2 == 0))
self.assertEqual(slot["preserve"], (i % 2 == 1))
def test_lora_count_truncates_slots(self):
node = self._node()
(out,) = node.build_options("simple", 2, **self._full_slot_kwargs())
self.assertEqual(len(out["slots"]), 2) # only 2 though 10 widgets set
def test_missing_widget_uses_slot_default(self):
# widgets not passed fall back to _SLOT_DEFAULTS (single source)
node = self._node()
(out,) = node.build_options("simple", 1)
self.assertEqual(out["slots"][0],
dict(lora_optimizer.LoRAOptimizerInline._SLOT_DEFAULTS))
def test_input_types_surface(self):
cls = lora_optimizer.LoRAInlineChainOptions
it = cls.INPUT_TYPES()
req = it["required"]
self.assertIn("settings_visibility", req)
self.assertIn("lora_count", req)
for i in (1, cls.MAX_LORAS):
for base in ("enabled", "strength", "model_strength",
"clip_strength", "conflict_mode", "key_filter",
"preserve"):
self.assertIn(f"{base}_{i}", req)
self.assertNotIn("optional", it) # pure widget node, no sockets
self.assertEqual(cls.FUNCTION, "build_options")
self.assertEqual(cls.RETURN_TYPES, ("LORA_CHAIN_OPTIONS",))
self.assertEqual(cls.RETURN_NAMES, ("chain_options",))
class TestRegistration(unittest.TestCase):
def test_inline_node_registered(self):
self.assertIn("LoRAOptimizerInline", lora_optimizer.NODE_CLASS_MAPPINGS)
self.assertIs(lora_optimizer.NODE_CLASS_MAPPINGS["LoRAOptimizerInline"],
lora_optimizer.LoRAOptimizerInline)
self.assertIn("LoRAOptimizerInline",
lora_optimizer.NODE_DISPLAY_NAME_MAPPINGS)
self.assertEqual(
lora_optimizer.NODE_DISPLAY_NAME_MAPPINGS["LoRAOptimizerInline"],
"LoRA Optimizer (Inline Chain)")
def test_options_node_registered(self):
self.assertIn("LoRAInlineChainOptions",
lora_optimizer.NODE_CLASS_MAPPINGS)
self.assertIs(
lora_optimizer.NODE_CLASS_MAPPINGS["LoRAInlineChainOptions"],
lora_optimizer.LoRAInlineChainOptions)
self.assertEqual(
lora_optimizer.NODE_DISPLAY_NAME_MAPPINGS["LoRAInlineChainOptions"],
"LoRA Inline Chain Options")
def _advanced_settings(**overrides):
"""Full OPTIMIZER_SETTINGS advanced-mode dict with deterministic,
test-friendly engine knobs (no merge cache, no compression SVD, CPU SVD).
The settings input is the node's supported surface for engine kwargs —
execute_inline deliberately does not forward loose **kwargs to
optimize_merge."""
settings = {
"mode": "advanced",
"auto_strength": "disabled", "auto_strength_floor": -1.0,
"optimization_mode": "per_prefix",
"sparsification": "disabled", "sparsification_density": 0.7,
"dare_dampening": 0.0,
"merge_strategy_override": "",
"merge_refinement": "none", "strategy_set": "full",
"normalize_keys": "disabled",
"architecture_preset": "dit", "decision_smoothing": 0.25,
"smooth_slerp_gate": False,
"star_eta": 100.0, "tame_layers": 0.0, "tame_threshold": 0.3,
"cache_patches": "disabled", "patch_compression": "disabled",
"svd_device": "cpu", "free_vram_between_passes": "disabled",
"vram_budget": 0.0,
}
settings.update(overrides)
return settings
class TestMultiLoraEndToEnd(unittest.TestCase):
"""Captured MULTI-LoRA chains through the REAL optimize_merge: two rank-4
adapters on the same 16x16 layer key, different loader strengths. The
single-LoRA virtual path is covered above; this closes the >= 2-LoRA gap
(full Pass 1 conflict analysis + Pass 2 group merge on virtual items)."""
KEY = "layer.weight"
@staticmethod
def _det_mats(seed_row):
"""Deterministic rank-2 up/down matrices (each linspace reshape has
rows/columns affine in the index, so the span is {1, arange} — rank 2;
deterministic on purpose, no randn luck)."""
up = torch.linspace(-1.0 + seed_row, 1.0 + seed_row, 16 * 4).reshape(16, 4)
down = torch.linspace(0.5 - seed_row, -0.5 + seed_row, 4 * 16).reshape(4, 16)
return up, down
def _run(self, chain, settings=None, output_strength=1.0, **exec_kwargs):
applied = {}
model = _pipeline_model(_chain_patches(*chain), applied)
# Snapshot the loader entries BEFORE execute_inline runs — comparing
# a post-run snapshot to the post-run state would be tautological and
# could never catch the node mutating its input model's patch lists.
orig_entries = {k: list(v) for k, v in model.patches.items()}
node = lora_optimizer.LoRAOptimizerInline()
node._get_model_keys = lambda m: {"alias_layer": self.KEY}
calls = []
with mock.patch.object(lora_optimizer.comfy.sd, "load_lora_for_models",
_realistic_load_lora_for_models(calls)):
result = node.execute_inline(model, output_strength=output_strength,
settings=settings, **exec_kwargs)
out = result["result"] if isinstance(result, dict) else result
return model, applied, calls, out, orig_entries
def test_two_lora_chain_weighted_sum_is_numerically_exact(self):
up_a, down_a = self._det_mats(0.0)
up_b, down_b = self._det_mats(0.25)
chain = (
(1.0, {self.KEY: _adapter(up=up_a, down=down_a)}),
(0.7, {self.KEY: _adapter(up=up_b, down=down_b)}),
)
# weighted_sum override makes the expected result exact:
# merged = 1.0 * (up_a @ down_a) + 0.7 * (up_b @ down_b)
settings = _advanced_settings(optimization_mode="global",
merge_strategy_override="weighted_sum")
model, applied, calls, out, orig_entries = self._run(
chain, settings=settings)
# (a) report opens with the chain fingerprints, engine report follows
report = out[2]
self.assertIn("#1: 1 keys, rank 4, loader strength 1.00", report)
self.assertIn("#2: 1 keys, rank 4, loader strength 0.70", report)
self.assertIn("chain lora #1", report)
# (b) merged patches re-applied via add_patches on a clone
self.assertEqual(calls, []) # no fast-path bypass
self.assertIn("patches", applied)
self.assertIsNot(applied["patcher"], model)
self.assertIs(out[0], applied["patcher"]) # patched clone returned
self.assertAlmostEqual(applied["strength"], 1.0)
# (c) the applied patch is the exact weighted combination
self.assertEqual(set(applied["patches"]), {self.KEY})
got = lora_optimizer._LoRAMergeBase._expand_patch_to_diff(
applied["patches"][self.KEY])
expected = up_a @ down_a + 0.7 * (up_b @ down_b)
self.assertTrue(torch.allclose(got, expected, atol=1e-5),
f"max err {(got - expected).abs().max().item()}")
# (d) the original input model still carries its loader entries
# (orig_entries snapshotted pre-run in _run, so this has teeth)
self.assertEqual(model.patches[self.KEY], orig_entries[self.KEY])
self.assertEqual(len(model.patches[self.KEY]), 2)
def test_output_strength_scales_at_apply_not_in_patches(self):
# Same exact-weighted-sum setup, but output_strength=0.5: the master
# volume must ride on add_patches ONCE and never ALSO be baked into
# the patch tensors (which would double-scale the merged result).
up_a, down_a = self._det_mats(0.0)
up_b, down_b = self._det_mats(0.25)
chain = (
(1.0, {self.KEY: _adapter(up=up_a, down=down_a)}),
(0.7, {self.KEY: _adapter(up=up_b, down=down_b)}),
)
settings = _advanced_settings(optimization_mode="global",
merge_strategy_override="weighted_sum")
model, applied, calls, out, orig_entries = self._run(
chain, settings=settings, output_strength=0.5)
self.assertEqual(calls, [])
self.assertAlmostEqual(applied["strength"], 0.5)
got = lora_optimizer._LoRAMergeBase._expand_patch_to_diff(
applied["patches"][self.KEY])
# UNscaled weighted sum — identical to the output_strength=1.0 case
expected = up_a @ down_a + 0.7 * (up_b @ down_b)
self.assertTrue(torch.allclose(got, expected, atol=1e-5),
f"max err {(got - expected).abs().max().item()}")
def test_none_alpha_captured_adapters_merge_without_crash(self):
# End-to-end guard: a chain of PEFT/diffusers LoRAs (alpha=None on the
# captured comfy adapters) must merge and re-apply without a
# TypeError, treating None alpha as scale 1.0.
up_a, down_a = self._det_mats(0.0)
up_b, down_b = self._det_mats(0.25)
chain = (
(1.0, {self.KEY: _adapter(up=up_a, down=down_a, alpha=None)}),
(0.7, {self.KEY: _adapter(up=up_b, down=down_b, alpha=None)}),
)
settings = _advanced_settings(optimization_mode="global",
merge_strategy_override="weighted_sum")
model, applied, calls, out, orig_entries = self._run(
chain, settings=settings)
self.assertEqual(calls, [])
self.assertIn("patches", applied)
got = lora_optimizer._LoRAMergeBase._expand_patch_to_diff(
applied["patches"][self.KEY])
self.assertTrue(torch.isfinite(got).all())
# alpha=None -> scale 1.0, so the exact weighted sum is unchanged
expected = up_a @ down_a + 0.7 * (up_b @ down_b)
self.assertTrue(torch.allclose(got, expected, atol=1e-5),
f"max err {(got - expected).abs().max().item()}")
def test_conflicting_loras_complete_through_conflict_analysis(self):
# Opposite-sign directions built deliberately: B is A negated with a
# small deterministic perturbation (near-antiparallel, cos ~ -1) so
# the per_prefix conflict path really engages without relying on a
# degenerate exactly-antiparallel corner.
up_a, down_a = self._det_mats(0.0)
down_b = down_a.clone()
down_b[0] += 0.05
chain = (
(1.0, {self.KEY: _adapter(up=up_a, down=down_a)}),
(0.7, {self.KEY: _adapter(up=-up_a, down=down_b)}),
)
model, applied, calls, out, orig_entries = self._run(
chain, settings=_advanced_settings())
self.assertEqual(calls, [])
report = out[2]
self.assertIn("#1: 1 keys, rank 4, loader strength 1.00", report)
self.assertIn("#2: 1 keys, rank 4, loader strength 0.70", report)
# merged patches exist and are finite/nonzero for the shared key
self.assertEqual(set(applied["patches"]), {self.KEY})
got = lora_optimizer._LoRAMergeBase._expand_patch_to_diff(
applied["patches"][self.KEY])
self.assertEqual(tuple(got.shape), (16, 16))
self.assertTrue(torch.isfinite(got).all())
self.assertGreater(got.abs().max().item(), 0.0)
# direction teeth: A (strength 1.0) beats the weaker near-antiparallel
# B (0.7), so any sane conflict resolution must land on A's side of
# the axis — positive cosine alignment with A's diff.
diff_a = (up_a @ down_a).flatten()
cos = torch.nn.functional.cosine_similarity(
got.flatten(), diff_a, dim=0).item()
self.assertGreater(cos, 0.0,
f"merged patch flipped to the weaker LoRA's "
f"direction (cos={cos:.4f})")
# original chain untouched on the input model (pre-run snapshot)
self.assertEqual(model.patches[self.KEY], orig_entries[self.KEY])
self.assertEqual(len(model.patches[self.KEY]), 2)
# Model-space target keys the loader chain leaves on the ModelPatcher. These
# still carry structural architecture markers even though they are NOT
# trainer-format LoRA keys, so _detect_architecture's structural heuristics
# resolve them.
_LTX_KEY = "diffusion_model.transformer_blocks.0.attn1.to_q.weight"
_ZIMAGE_KEY = "diffusion_model.layers.0.attention.to_q.weight"
# SDXL UNet cross-attention: has transformer_blocks + attn1 (which naively
# looks like LTX) but lives under input_blocks -> the LTX branch's UNet-block
# exclusion must keep it out of 'ltx' and land it on sd15 -> sd_unet.
_SDXL_UNET_KEY = ("diffusion_model.input_blocks.4.1.transformer_blocks.0"
".attn1.to_q.weight")
class TestVirtualArchDetection(unittest.TestCase):
"""_normalize_stack must fall back to detecting architecture from the
MODEL-SPACE keys of virtual (_precomputed_diffs) items when no file-based
item resolves. Inline-captured chains are 100% virtual, so without this
fallback every inline merge resolved to sd_unet (wrong for DiT models)."""
def _virtual_stack(self, model_key):
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {model_key: _adapter()})))
return lora_optimizer.LoRAOptimizerInline._chain_groups_to_stack(
mg, [], [_slot()], "simple")
def test_ltx_model_keys_detect_ltx(self):
node = lora_optimizer.LoRAOptimizerInline()
stack = self._virtual_stack(_LTX_KEY)
self.assertTrue(all(i["_precomputed_diffs"] for i in stack))
node._normalize_stack(stack)
self.assertEqual(node._detected_arch, "ltx")
def test_zimage_model_keys_detect_zimage(self):
node = lora_optimizer.LoRAOptimizerInline()
stack = self._virtual_stack(_ZIMAGE_KEY)
node._normalize_stack(stack)
self.assertEqual(node._detected_arch, "zimage")
def test_tuple_keys_do_not_crash_and_still_detect(self):
# Fused-QKV captures key the virtual dict by (str_key, offset) TUPLES.
# _detect_architecture indexes k.lower()/'in k', which would crash on a
# tuple -> the fallback must stringify the key view first.
off = (0, 0, 4)
patches = {_LTX_KEY: [_entry(0.7, _adapter(), offset=off)]}
groups = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(patches)
self.assertEqual(list(groups[0]["entries"]), [(_LTX_KEY, off)])
stack = lora_optimizer.LoRAOptimizerInline._chain_groups_to_stack(
groups, [], [_slot()], "simple")
node = lora_optimizer.LoRAOptimizerInline()
node._normalize_stack(stack) # must not raise on the tuple key
self.assertEqual(node._detected_arch, "ltx")
# the real lora dict is untouched — its key is still the TUPLE
self.assertIn((_LTX_KEY, off), stack[0]["lora"])
def test_file_item_detection_not_overridden_by_virtual_fallback(self):
# Regression guard: a file-based item with recognizable trainer keys
# resolves first (break), so the virtual LTX item never reaches the
# fallback. Proves the fallback fires ONLY when file detection fails.
file_item = {
"name": "sdxl_style.safetensors",
"lora": {"lora_te1_text_model_encoder_layers_0_self_attn_q_proj"
".lora_up.weight": torch.zeros(4, 4)},
"strength": 1.0,
}
virtual = self._virtual_stack(_LTX_KEY)[0]
node = lora_optimizer.LoRAOptimizerInline()
node._normalize_stack([file_item, virtual])
self.assertEqual(node._detected_arch, "sdxl") # file wins, not 'ltx'
def test_sdxl_unet_virtual_keys_resolve_to_sd_unet_not_ltx(self):
# No false LTX match: SDXL UNet cross-attn keys carry transformer_blocks
# + attn1 but under input_blocks -> sd15 -> sd_unet preset.
node = lora_optimizer.LoRAOptimizerInline()
stack = self._virtual_stack(_SDXL_UNET_KEY)
node._normalize_stack(stack)
self.assertEqual(node._detected_arch, "sd15")
self.assertNotEqual(node._detected_arch, "ltx")
key, _ = lora_optimizer._resolve_arch_preset("auto", node._detected_arch)
self.assertEqual(key, "sd_unet")
def test_resolve_arch_preset_ltx_and_zimage_map_to_dit(self):
# Prove the preset now resolves correctly once detection works.
self.assertEqual(lora_optimizer._resolve_arch_preset("auto", "ltx")[0],
"dit")
self.assertEqual(
lora_optimizer._resolve_arch_preset("auto", "zimage")[0], "dit")
class TestInlineFingerprintArch(unittest.TestCase):
"""execute_inline must detect the architecture from the captured
model-space keys and surface it in the fingerprint, instead of always
warning 'architecture: unknown (inline capture)'."""
def _model(self, patches):
m = _FakePatcher(patches)
m.model = types.SimpleNamespace()
return m
@staticmethod
def _out(result):
return result["result"] if isinstance(result, dict) else result
def test_ltx_chain_reports_detected_arch(self):
model = self._model(_chain_patches((0.8, {_LTX_KEY: _adapter()})))
node = lora_optimizer.LoRAOptimizerInline()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "engine report",
None, None)
report = self._out(node.execute_inline(model, output_strength=1.0))[2]
self.assertIn("architecture: ltx", report)
self.assertIn("detected from captured keys", report)
self.assertNotIn("architecture: unknown (inline capture)", report)
def test_unknown_chain_still_warns(self):
# Regression net: an unrecognizable key with no pinned preset still
# emits the "set a preset" guidance.
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = lora_optimizer.LoRAOptimizerInline()
node.optimize_merge = lambda m, s, o, **kw: (m, None, "r", None, None)
report = self._out(node.execute_inline(model, output_strength=1.0))[2]
self.assertIn("architecture: unknown (inline capture)", report)
self.assertIn("architecture_preset", report)
class TestInlineRankReporting(unittest.TestCase):
"""The inline analysis path must read each captured adapter's TRUE rank
(mat_down.shape[0]) instead of counting every precomputed diff as rank 1.
A rank-1 undercount makes avg_rank -> sum_rank -> compress_rank floor to
64, over-compressing higher-rank captured LoRAs."""
KEY = "layer.weight"
# ---- unit: the _payload_rank helper, now on the shared base class ----
def test_payload_rank_reads_adapter_rank(self):
got = lora_optimizer._LoRAMergeBase._payload_rank(
_adapter(rank=32, out_dim=16, in_dim=16))
self.assertEqual(got, 32)
def test_payload_rank_none_for_dense_diff_tuple(self):
self.assertIsNone(lora_optimizer._LoRAMergeBase._payload_rank(
("diff", (torch.randn(8, 8),))))
def test_payload_rank_none_for_bare_tensor(self):
self.assertIsNone(
lora_optimizer._LoRAMergeBase._payload_rank(torch.randn(8, 8)))
def test_payload_rank_none_for_lokr(self):
lokr = lora_optimizer.LoKrAdapter(
loaded_keys=set(),
weights=(torch.randn(4, 4), torch.randn(2, 2), None, None,
None, None, None, None))
self.assertIsNone(lora_optimizer._LoRAMergeBase._payload_rank(lokr))
def test_helper_not_left_on_inline_subclass(self):
# Moved to the base, so the subclass must NOT redeclare it (DRY: one
# implementation shared by the fingerprint AND _prepare_group_diffs).
self.assertNotIn("_payload_rank",
vars(lora_optimizer.LoRAOptimizerInline))
self.assertIn("_payload_rank",
vars(lora_optimizer._LoRAMergeBase))
# ---- _prepare_group_diffs rank accounting on virtual items ----
def _rank_sums(self, payloads, dim=16):
"""Run _prepare_group_diffs over a virtual (_precomputed_diffs) stack
of the given per-LoRA payloads on one shared model key; return the
resolved rank_sums dict."""
model = _pipeline_model({}, {}, dim=dim) # .model has layer.weight
active = [{"name": f"chain lora #{i + 1}", "strength": 1.0,
"_precomputed_diffs": True, "lora": {self.KEY: p}}
for i, p in enumerate(payloads)]
target_group = {"target_key": self.KEY, "is_clip": False,
"aliases": [self.KEY], "label_prefix": self.KEY}
node = lora_optimizer.LoRAOptimizerInline()
prepared = node._prepare_group_diffs(
target_group, active, model, None, torch.device("cpu"))
return prepared["rank_sums"]
def test_prepare_group_diffs_reads_true_adapter_rank(self):
rank_sums = self._rank_sums([_adapter(rank=32, out_dim=16, in_dim=16)])
self.assertEqual(rank_sums[0], 32) # was 1 before the fix
def test_prepare_group_diffs_dense_diff_stays_rank_1(self):
# A formula-submerge-style dense virtual payload has no meaningful rank
# -> _payload_rank is None -> fall back to += 1 (unchanged behavior).
rank_sums = self._rank_sums([("diff", (torch.randn(16, 16),))])
self.assertEqual(rank_sums[0], 1)
def test_prepare_group_diffs_bare_tensor_stays_rank_1(self):
rank_sums = self._rank_sums([torch.randn(16, 16)])
self.assertEqual(rank_sums[0], 1)
# ---- end-to-end: avg_rank -> sum_rank -> compress_rank ----
def _run_inline(self, chain, settings):
applied = {}
model = _pipeline_model(_chain_patches(*chain), applied)
node = lora_optimizer.LoRAOptimizerInline()
node._get_model_keys = lambda m: {"alias_layer": self.KEY}
calls = []
with mock.patch.object(lora_optimizer.comfy.sd, "load_lora_for_models",
_realistic_load_lora_for_models(calls)):
result = node.execute_inline(model, output_strength=1.0,
settings=settings)
return result["result"] if isinstance(result, dict) else result
def test_report_shows_true_avg_rank_not_one(self):
chain = (
(1.0, {self.KEY: _adapter(rank=32, out_dim=16, in_dim=16)}),
(0.7, {self.KEY: _adapter(rank=32, out_dim=16, in_dim=16)}),
)
settings = _advanced_settings(optimization_mode="global",
merge_strategy_override="weighted_sum")
report = self._run_inline(chain, settings)[2]
self.assertIn("Avg rank: 32", report)
self.assertNotIn("Avg rank: 1\n", report) # the old always-1 bug
def test_compress_rank_not_floored_for_high_rank_chain(self):
# The user's actual regression: rank-64 + rank-128 captured LoRAs.
# True sum_rank = 64 + 128 = 192 -> compress_rank = max(192, 64) = 192.
# The bug reported avg_rank 1 each -> sum_rank 2 -> floored to 64,
# over-compressing the rank-128 LoRA.
chain = (
(1.0, {self.KEY: _adapter(rank=64, out_dim=16, in_dim=16)}),
(0.7, {self.KEY: _adapter(rank=128, out_dim=16, in_dim=16)}),
)
settings = _advanced_settings(optimization_mode="global",
merge_strategy_override="weighted_sum",
patch_compression="smart")
lora_data = self._run_inline(chain, settings)[4]
self.assertEqual(lora_data["sum_rank"], 192) # NOT floored to 64
def _autotuner_settings(**overrides):
"""Complete OPTIMIZER_SETTINGS autotuner-mode dict, mirroring what the
AutoTuner Settings node (build_settings) emits — with cheap, deterministic
knobs. execute_inline reads the bracket-keyed fields to build the auto_tune
call, so a partial dict would KeyError before auto_tune is even reached."""
settings = {
"mode": "autotuner",
"top_n": 1,
"scoring_svd": "disabled",
"scoring_device": "cpu",
"scoring_speed": "full",
"scoring_formula": "v2",
"output_mode": "merge",
"smooth_slerp_gate": False,
"normalize_keys": "enabled", # must be overridden to disabled
"architecture_preset": "auto",
"auto_strength_floor": -1.0,
"decision_smoothing": 0.25,
"vram_budget": 0.0,
"cache_patches": "disabled",
"star_eta": 100.0,
"tame_layers": 0.0,
"tame_threshold": 0.3,
"diff_cache_mode": "disabled",
"diff_cache_ram_pct": 0.5,
"community_cache": "disabled",
"evaluator": None,
"memory_mode": "disabled",
"selection": 1,
"record_dataset": "disabled",
}
settings.update(overrides)
return settings
class TestInlineAutoTunerDelegation(unittest.TestCase):
"""Commit B: mode='autotuner' settings now delegate to the real AutoTuner
on the inline-prepared STRIPPED model/clip + VIRTUAL stack, forcing
normalize_keys off, instead of falling back to optimizer defaults."""
def _node(self):
return lora_optimizer.LoRAOptimizerInline()
def _model(self, patches):
m = _FakePatcher(patches)
m.model = types.SimpleNamespace()
return m
@staticmethod
def _out(result):
return result["result"] if isinstance(result, dict) else result
def _stub_delegate(self, node, seen):
"""Pre-seed node._autotuner_delegate with a recorder so execute_inline
reuses it instead of building a real LoRAAutoTuner."""
def fake_auto_tune(model, stack, output_strength, **kw):
seen["model"] = model
seen["stack"] = stack
seen["output_strength"] = output_strength
seen.update(kw)
return (model, kw.get("clip"), "tuner report",
"analysis report", {"td": 1}, {"ld": 2})
node._autotuner_delegate = types.SimpleNamespace(auto_tune=fake_auto_tune)
def test_delegates_with_virtual_stack_and_stripped_model(self):
model = self._model(_chain_patches(
(0.8, {"a": _adapter()}), (0.5, {"a": _adapter()})))
node = self._node()
node.optimize_merge = lambda *a, **k: self.fail(
"optimize_merge must not run in autotuner mode")
seen = {}
self._stub_delegate(node, seen)
node.execute_inline(model, output_strength=1.0,
settings=_autotuner_settings())
# VIRTUAL stack (precomputed diffs) ...
self.assertEqual(len(seen["stack"]), 2)
self.assertTrue(all(i["_precomputed_diffs"] for i in seen["stack"]))
# ... STRIPPED model (patches removed, not the raw input) ...
self.assertEqual(seen["model"].patches, {})
self.assertIsNot(seen["model"], model)
# ... normalize_keys forced off despite settings "enabled"
self.assertEqual(seen["normalize_keys"], "disabled")
self.assertAlmostEqual(seen["output_strength"], 1.0)
def test_return_is_inline_5_tuple_with_fingerprint_prepended(self):
model = self._model(_chain_patches((0.8, {"a": _adapter(), "b": _adapter()})))
node = self._node()
seen = {}
self._stub_delegate(node, seen)
out = self._out(node.execute_inline(model, output_strength=1.0,
settings=_autotuner_settings()))
self.assertEqual(len(out), 5) # inline 5-tuple, not auto_tune's 6
report = out[2]
self.assertIn("Detected loader chain", report) # fingerprint
self.assertLess(report.index("Detected loader chain"),
report.index("tuner report")) # prepended
self.assertIn("ANALYSIS REPORT", report) # analysis folded in
self.assertIn("analysis report", report)
self.assertEqual(out[3], {"td": 1}) # tuner_data
self.assertEqual(out[4], {"ld": 2}) # lora_data
def test_forwards_settings_fields_to_auto_tune(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
seen = {}
self._stub_delegate(node, seen)
node.execute_inline(model, output_strength=1.0,
settings=_autotuner_settings(
top_n=3, memory_mode="auto",
community_cache="upload_and_download",
architecture_preset="dit", cache_patches="enabled"))
self.assertEqual(seen["top_n"], 3)
self.assertEqual(seen["memory_mode"], "auto")
self.assertEqual(seen["community_cache"], "upload_and_download")
self.assertEqual(seen["architecture_preset"], "dit")
self.assertEqual(seen["cache_patches"], "enabled") # forwarded from settings
def test_passes_stripped_clip_original_untouched(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
clip = _FakeCLIP(_FakePatcher(_chain_patches((0.6, {"te.a": _adapter()}))))
node = self._node()
seen = {}
self._stub_delegate(node, seen)
node.execute_inline(model, output_strength=1.0, clip=clip,
settings=_autotuner_settings())
self.assertIsNot(seen["clip"], clip) # stripped clone
self.assertEqual(seen["clip"].patcher.patches, {})
self.assertEqual(len(clip.patcher.patches["te.a"]), 1) # original intact
def test_old_not_supported_message_gone(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
seen = {}
self._stub_delegate(node, seen)
report = self._out(node.execute_inline(
model, output_strength=1.0, settings=_autotuner_settings()))[2]
self.assertNotIn("AutoTuner settings are not supported inline", report)
self.assertNotIn("using optimizer defaults", report)
def test_pinned_preset_suppresses_arch_unknown_warning(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
seen = {}
self._stub_delegate(node, seen)
report = self._out(node.execute_inline(
model, output_strength=1.0,
settings=_autotuner_settings(architecture_preset="dit")))[2]
self.assertNotIn("architecture: unknown", report)
def test_delegate_created_lazily_and_reused(self):
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
self.assertFalse(hasattr(node, "_autotuner_delegate"))
with mock.patch.object(lora_optimizer.LoRAAutoTuner, "auto_tune",
return_value=(model, None, "r", "", None, None)) as m:
node.execute_inline(model, output_strength=1.0,
settings=_autotuner_settings())
delegate = node._autotuner_delegate
self.assertIsInstance(delegate, lora_optimizer.LoRAAutoTuner)
node.execute_inline(model, output_strength=1.0,
settings=_autotuner_settings())
self.assertIs(node._autotuner_delegate, delegate) # reused, not recreated
self.assertEqual(m.call_count, 2)
def test_advanced_mode_still_uses_optimize_merge(self):
# Regression: advanced mode is untouched — it must still route through
# optimize_merge (not the AutoTuner) with normalize_keys pinned off.
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
seen = {}
def fake_merge(m, stack, output_strength, **kw):
seen.update(kw)
seen["stack"] = stack
return (m, kw.get("clip"), "engine report", None, None)
node.optimize_merge = fake_merge
node._autotuner_delegate = types.SimpleNamespace(
auto_tune=lambda *a, **k: self.fail("auto_tune called in advanced mode"))
node.execute_inline(model, output_strength=1.0, settings=_advanced_settings())
self.assertEqual(seen["normalize_keys"], "disabled")
self.assertIn("stack", seen)
def test_no_settings_path_still_uses_optimize_merge(self):
# Regression: no settings -> plain optimize_merge, no AutoTuner.
model = self._model(_chain_patches((0.8, {"a": _adapter()})))
node = self._node()
seen = {}
def fake_merge(m, stack, output_strength, **kw):
seen.update(kw)
return (m, kw.get("clip"), "engine report", None, None)
node.optimize_merge = fake_merge
node._autotuner_delegate = types.SimpleNamespace(
auto_tune=lambda *a, **k: self.fail("auto_tune called with no settings"))
node.execute_inline(model, output_strength=1.0)
self.assertEqual(seen["normalize_keys"], "disabled")
class TestInlineAutoTunerEndToEnd(unittest.TestCase):
"""The REAL AutoTuner runs a tiny search over a 2-captured-LoRA chain and
returns a merged, re-applied model — proving auto_tune tolerates virtual
(_precomputed_diffs) items end to end, with no file reload."""
KEY = "layer.weight"
def test_two_lora_chain_autotunes_and_reapplies(self):
up_a = torch.linspace(-1.0, 1.0, 16 * 4).reshape(16, 4)
down_a = torch.linspace(0.5, -0.5, 4 * 16).reshape(4, 16)
up_b = torch.linspace(-0.8, 1.2, 16 * 4).reshape(16, 4)
down_b = torch.linspace(0.3, -0.7, 4 * 16).reshape(4, 16)
chain = (
(1.0, {self.KEY: _adapter(up=up_a, down=down_a)}),
(0.7, {self.KEY: _adapter(up=up_b, down=down_b)}),
)
applied = {}
model = _pipeline_model(_chain_patches(*chain), applied)
orig_entries = {k: list(v) for k, v in model.patches.items()}
node = lora_optimizer.LoRAOptimizerInline()
node._get_model_keys = lambda m: {"alias_layer": self.KEY}
# The merge runs on the delegate (a separate LoRAAutoTuner), so the
# key-alias override the fake model needs must live there too. In
# production both share the real _get_model_keys against a real model.
node._autotuner_delegate = lora_optimizer.LoRAAutoTuner()
node._autotuner_delegate._get_model_keys = lambda m: {"alias_layer": self.KEY}
calls = []
with mock.patch.object(lora_optimizer.comfy.sd, "load_lora_for_models",
_realistic_load_lora_for_models(calls)):
result = node.execute_inline(
model, output_strength=1.0,
settings=_autotuner_settings(architecture_preset="dit"))
out = result["result"] if isinstance(result, dict) else result
# a real AutoTuner ran the search
self.assertIsInstance(node._autotuner_delegate,
lora_optimizer.LoRAAutoTuner)
self.assertEqual(calls, []) # no single-LoRA fast-path
# merged patch re-applied on a clone for the shared key
self.assertIn("patches", applied)
self.assertEqual(set(applied["patches"]), {self.KEY})
got = lora_optimizer._LoRAMergeBase._expand_patch_to_diff(
applied["patches"][self.KEY])
self.assertEqual(tuple(got.shape), (16, 16))
self.assertTrue(torch.isfinite(got).all())
self.assertGreater(got.abs().max().item(), 0.0)
# inline 5-tuple with the fingerprint prepended
self.assertEqual(len(out), 5)
self.assertIn("Detected loader chain", out[2])
# input model untouched (pre-run snapshot)
self.assertEqual(model.patches[self.KEY], orig_entries[self.KEY])
class TestCapturedContentHash(unittest.TestCase):
"""Commit A: _lora_content_hash must produce a STABLE 16-hex content hash
for captured chain items (adapter objects, tuple keys, diff tuples),
keyed on the factor VALUES not object identity, so inline chains get a
persistent memory/community identity across sessions and machines. The
file-based path and plain-tensor path stay byte-identical."""
def _hash(self, item):
# Captured items have no file on disk -> force the in-memory fallback.
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
return lora_optimizer.LoRAAutoTuner._lora_content_hash(item)
def test_captured_adapter_item_hashes_to_stable_16hex(self):
up, down = torch.randn(8, 4), torch.randn(4, 8)
h = self._hash({"name": "chain lora #1 [x]",
"lora": {"k": _adapter(up=up, down=down)}})
self.assertIsNotNone(h)
self.assertEqual(len(h), 16)
int(h, 16) # valid hex
def test_hash_identical_for_fresh_adapters_with_same_values(self):
# Two DIFFERENT adapter OBJECTS holding the SAME tensor values must
# hash identically -> content-based, not identity-based. Pre-fix this
# failed: adapters have no .detach, so repr(v) (memory address) was
# hashed and differed per object.
up, down = torch.randn(8, 4), torch.randn(4, 8)
a1 = _adapter(up=up.clone(), down=down.clone())
a2 = _adapter(up=up.clone(), down=down.clone())
self.assertIsNot(a1, a2)
h1 = self._hash({"name": "chain lora #1 [aaa]", "lora": {"k": a1}})
h2 = self._hash({"name": "chain lora #1 [bbb]", "lora": {"k": a2}})
self.assertEqual(h1, h2)
def test_hash_differs_when_a_weight_value_changes(self):
up, down = torch.randn(8, 4), torch.randn(4, 8)
h1 = self._hash({"name": "n",
"lora": {"k": _adapter(up=up.clone(), down=down.clone())}})
down2 = down.clone()
down2[0, 0] += 1.0
h2 = self._hash({"name": "n",
"lora": {"k": _adapter(up=up.clone(), down=down2)}})
self.assertNotEqual(h1, h2)
def test_hash_stable_when_alpha_changes_is_reflected(self):
# alpha is a non-tensor weight element (repr-hashed) -> a different
# alpha must change the hash.
up, down = torch.randn(8, 4), torch.randn(4, 8)
h1 = self._hash({"name": "n",
"lora": {"k": _adapter(up=up.clone(), down=down.clone(),
alpha=4.0)}})
h2 = self._hash({"name": "n",
"lora": {"k": _adapter(up=up.clone(), down=down.clone(),
alpha=2.0)}})
self.assertNotEqual(h1, h2)
def test_tuple_keys_hash_without_crashing_and_stably(self):
# Fused-QKV captures key the virtual dict by (str_key, offset) TUPLES.
# Pre-fix k.encode() crashed on a tuple; str(k) must be used.
off = (0, 0, 4)
up, down = torch.randn(8, 4), torch.randn(4, 8)
item1 = {"name": "n",
"lora": {("k", off): _adapter(up=up.clone(), down=down.clone())}}
item2 = {"name": "n",
"lora": {("k", off): _adapter(up=up.clone(), down=down.clone())}}
h1 = self._hash(item1) # must not raise
self.assertEqual(len(h1), 16)
self.assertEqual(h1, self._hash(item2))
def test_diff_tuple_payload_hashes_stably_and_value_sensitively(self):
t = torch.randn(4, 4)
item1 = {"name": "n", "lora": {"k": ("diff", (t.clone(),))}}
item2 = {"name": "n", "lora": {"k": ("diff", (t.clone(),))}}
h1 = self._hash(item1)
self.assertEqual(len(h1), 16)
self.assertEqual(h1, self._hash(item2))
t2 = t.clone()
t2[0, 0] += 1.0
self.assertNotEqual(h1, self._hash({"name": "n",
"lora": {"k": ("diff", (t2,))}}))
def test_plain_tensor_dict_still_hashes_and_is_stable(self):
# Regression: LoRAExtractFromModel-style plain-tensor dicts still hash.
sd = {"blk.lora_up.weight": torch.ones(4, 2),
"blk.lora_down.weight": torch.ones(2, 4) * 0.5}
h1 = self._hash({"name": "<x>", "lora": sd})
self.assertIsNotNone(h1)
self.assertEqual(len(h1), 16)
self.assertEqual(h1, self._hash({"name": "<x>", "lora": dict(sd)}))
def test_plain_tensor_hash_byte_identical_to_prefix_reference(self):
# The plain-tensor code path must be UNCHANGED (existing
# LoRAExtractFromModel community entries must not shift). Reproduce the
# exact pre-fix per-key hashing and assert equality.
import hashlib
sd = {"b.up": torch.ones(4, 2), "b.down": torch.ones(2, 4) * 0.5}
h = hashlib.sha256()
for k in sorted(sd.keys()):
v = sd[k]
h.update(k.encode())
t = v.detach().to("cpu", torch.float32).contiguous()
h.update(str(tuple(t.shape)).encode())
h.update(t.numpy().tobytes())
expected = h.hexdigest()[:16]
self.assertEqual(self._hash({"name": "<x>", "lora": sd}), expected)
def test_two_different_captured_loras_differ(self):
a = {"name": "n", "lora": {"k": _adapter(up=torch.ones(8, 4),
down=torch.ones(4, 8))}}
b = {"name": "n", "lora": {"k": _adapter(up=torch.zeros(8, 4),
down=torch.ones(4, 8))}}
self.assertNotEqual(self._hash(a), self._hash(b))
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink