nodes/ethanfel-ComfyUI-LoRA-Optimizer
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
ca2b0e912548831207f233d6ff900e1f041ad987
ethanfel-ComfyUI-LoRA-Optim…/tests/test_inline_optimizer.py
T
EthanfelandClaude Fable 5 ca2b0e9125 fix(inline): conservative order-independent stamp matching to prevent dataset pollution 
Important 2: the old _resolve_stamp_names matched stamp[i]->group[i] by
absolute index, guarded only by a strength check. It mis-attributes when
same-strength disjoint groups reconstruct in a different order, or when an
unstamped loader precedes stamped ones and a strength collision passes tol.
A wrong _resolved_file_name drives the persistent identity, so the run would
write one file's analysis/memory/community stats under another file's key —
polluting the file-based dataset and degrading future Stack runs of that file
(the merge output itself is unaffected: it runs on the real captured diffs).

Now match by UNIQUE strength, order-independent: name group g iff (i) g's
strength is unique among groups within tol, (ii) exactly one stamp matches
g's strength within tol, and (iii) that stamp's strength is unique among
stamps within tol. Any collision/ambiguity leaves g unnamed (chain lora #N +
captured identity) — better anonymous than a wrong file identity. The
uniqueness-within-tol test also closes the 0.8/0.80001 false-positive window.

Pinned by TestInlineConservativeMatching (6) + TestInlineAmbiguousNoPollution
(1). Distinct-strength naming still works (existing matcher/read-path tests
green). Full suite 509; version unchanged.

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

2527 lines
120 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 os
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"))
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")
def test_cache_patches_pinned_disabled_for_delegate(self):
# cache_patches is NOT forwarded from settings — the delegate's in-node
# cache can never usefully hit for inline and reintroduces id()-reuse
# staleness, so it is pinned off regardless of the settings value.
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(cache_patches="enabled"))
self.assertEqual(seen["cache_patches"], "disabled")
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))
class TestCapturedPersistentIdentity(unittest.TestCase):
"""Critical #1: the persistent memory + analysis-cache keys must derive
from captured CONTENT (session-stable), not the per-session salted item
name — otherwise inline memory saves under a fresh key every ComfyUI
session and never reads back. File-based items are unchanged."""
KEY = "layer.weight"
def _captured_item(self, salt, up, down):
# One session's virtual stack item: identical content, session salt.
mg = lora_optimizer._LoRAMergeBase._reconstruct_chain_groups(
_chain_patches((0.8, {"k": _adapter(up=up.clone(), down=down.clone())})))
return lora_optimizer.LoRAOptimizerInline._chain_groups_to_stack(
mg, [], [_slot()], "simple", name_salt=f" [{salt}]")[0]
def test_persistent_key_stable_across_sessions_for_captured(self):
up, down = torch.randn(8, 4), torch.randn(4, 8)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
s1 = self._captured_item("aaaaaaaa", up, down)
s2 = self._captured_item("bbbbbbbb", up, down)
self.assertNotEqual(s1["name"], s2["name"]) # salted names differ
k1 = lora_optimizer.LoRAAutoTuner._persistent_lora_key(s1)
k2 = lora_optimizer.LoRAAutoTuner._persistent_lora_key(s2)
self.assertTrue(k1.startswith("captured:"))
self.assertEqual(k1, k2) # content identity is session-stable
def test_persistent_key_is_name_for_file_items(self):
# Regression: file-based items (no _precomputed_diffs) key on name.
item = {"name": "style.safetensors", "strength": 1.0, "lora": {}}
self.assertEqual(
lora_optimizer.LoRAAutoTuner._persistent_lora_key(item),
"style.safetensors")
def test_identity_hash_stable_across_sessions_for_captured(self):
# The per-LoRA/pair analysis DISK cache key (_lora_identity_hash) must
# also be session-stable for captured items.
up, down = torch.randn(8, 4), torch.randn(4, 8)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
h1 = lora_optimizer.LoRAAutoTuner._lora_identity_hash(
self._captured_item("aaaaaaaa", up, down))
h2 = lora_optimizer.LoRAAutoTuner._lora_identity_hash(
self._captured_item("bbbbbbbb", up, down))
self.assertEqual(len(h1), 16)
self.assertEqual(h1, h2)
def test_identity_hash_unchanged_for_file_items(self):
# Regression: file item identity hash == the pre-fix name+0+0 formula.
import hashlib, json
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
item = {"name": "x.safetensors", "strength": 1.0, "lora": {}}
got = lora_optimizer.LoRAAutoTuner._lora_identity_hash(item)
expected = hashlib.sha256(json.dumps(
("x.safetensors", 0, 0), separators=(",", ":")).encode()).hexdigest()[:16]
self.assertEqual(got, expected)
def test_names_only_hash_stable_across_sessions_for_captured(self):
# The whole-stack analysis cache key must be session-stable too.
up, down = torch.randn(8, 4), torch.randn(4, 8)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
s1 = [self._captured_item("aaaaaaaa", up, down)]
s2 = [self._captured_item("bbbbbbbb", up, down)]
h1, _ = lora_optimizer.LoRAAutoTuner._compute_names_only_hash(s1)
h2, _ = lora_optimizer.LoRAAutoTuner._compute_names_only_hash(s2)
self.assertEqual(h1, h2)
def test_per_lora_merge_signature_stable_across_sessions_for_captured(self):
# Feeds the memory settings_hash (part of the memory file path) — must
# not vary by salt or cross-session hits break.
up, down = torch.randn(8, 4), torch.randn(4, 8)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
sig1 = lora_optimizer.LoRAAutoTuner._per_lora_merge_signature(
[self._captured_item("aaaaaaaa", up, down)])
sig2 = lora_optimizer.LoRAAutoTuner._per_lora_merge_signature(
[self._captured_item("bbbbbbbb", up, down)])
self.assertEqual(sig1, sig2)
def test_memo_content_hash_computed_once(self):
# Memoized on the item dict so re-hashing the (expensive) captured
# factors doesn't happen per candidate / per memory+community lookup.
up, down = torch.randn(8, 4), torch.randn(4, 8)
item = self._captured_item("aaaaaaaa", up, down)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
with mock.patch.object(
lora_optimizer.LoRAAutoTuner, "_lora_content_hash",
wraps=lora_optimizer.LoRAAutoTuner._lora_content_hash) as spy:
a = lora_optimizer.LoRAAutoTuner._persistent_lora_key(item)
b = lora_optimizer.LoRAAutoTuner._memo_content_hash(item)
c = lora_optimizer.LoRAAutoTuner._persistent_lora_key(item)
self.assertEqual(a, c)
self.assertEqual(a, f"captured:{b}")
self.assertEqual(spy.call_count, 1) # computed exactly once
class TestInlineMemoryCrossSession(unittest.TestCase):
"""The headline: a memory entry saved in 'session 1' (one salt) is found
in 'session 2' (different salt, identical captured content) — driving the
REAL auto_tune through execute_inline. Fails on name-based keys (no hit),
passes on content-based keys."""
KEY = "layer.weight"
@staticmethod
def _chain():
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)
return (
(1.0, {TestInlineMemoryCrossSession.KEY: _adapter(up=up_a, down=down_a)}),
(0.7, {TestInlineMemoryCrossSession.KEY: _adapter(up=up_b, down=down_b)}),
)
def _run_session(self, salt_uuid):
applied = {}
model = _pipeline_model(_chain_patches(*self._chain()), applied)
model.patches_uuid = salt_uuid # drives the per-session name salt
node = lora_optimizer.LoRAOptimizerInline()
node._get_model_keys = lambda m: {"alias_layer": self.KEY}
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",
memory_mode="auto", top_n=1))
return result["result"] if isinstance(result, dict) else result
def test_memory_saved_in_session1_is_found_in_session2(self):
import tempfile
with tempfile.TemporaryDirectory() as memdir:
# captured content has no file -> force in-memory content hashing
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", memdir), \
mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
out1 = self._run_session(uuid.uuid4())
self.assertNotIn("MEMORY HIT", out1[2]) # session 1: full sweep
out2 = self._run_session(uuid.uuid4()) # different salt
self.assertIn("MEMORY HIT", out2[2]) # session 2: cross-session hit
self.assertEqual(len(out2), 5)
class _AttachHost:
"""Minimal object exposing the ModelPatcher attachment API."""
def __init__(self):
self._att = {}
def get_attachment(self, key):
return self._att.get(key, None)
def set_attachments(self, key, val):
self._att[key] = val
class _ClipWrap:
"""CLIP-like wrapper whose attachments live on its .patcher (real CLIP
objects don't expose get/set_attachments — only clip.patcher does)."""
def __init__(self):
self.patcher = _AttachHost()
class TestLoraNameStampInstaller(unittest.TestCase):
"""Commit 1: _install_lora_name_stamp wraps a LoraLoader-like class so
every load_lora call appends {name, strength_model, strength_clip} to the
'loraopt_chain_names' attachment on the returned model (and clip)."""
ATTACH = "loraopt_chain_names"
def _fresh_loader_cls(self):
class _L:
def load_lora(self, model, clip, lora_name,
strength_model, strength_clip):
return (model, clip)
return _L
def test_default_target_never_raises(self):
# With/without comfy's top-level `nodes` importable, the no-arg call
# must degrade gracefully (never raise at package import).
try:
result = lora_optimizer._install_lora_name_stamp()
except Exception as e: # pragma: no cover
self.fail(f"installer raised: {e}")
self.assertIn(result, (True, False))
def test_entries_accumulate_in_chain_order(self):
L = self._fresh_loader_cls()
self.assertTrue(lora_optimizer._install_lora_name_stamp(L))
inst = L()
m = _AttachHost()
inst.load_lora(m, None, "styles/a.safetensors", 0.8, 0.7)
inst.load_lora(m, None, "chars/b.safetensors", 0.5, 0.5)
stamps = m.get_attachment(self.ATTACH)
self.assertEqual([s["name"] for s in stamps],
["styles/a.safetensors", "chars/b.safetensors"])
self.assertEqual(stamps[0]["strength_model"], 0.8)
self.assertEqual(stamps[0]["strength_clip"], 0.7)
self.assertEqual(stamps[1]["strength_model"], 0.5)
def test_builds_new_list_never_mutates_upstream(self):
# The accumulator MUST build a NEW list each call: ModelPatcher.clone
# copies the attachment list by REFERENCE, so appending in place would
# corrupt the upstream (pre-clone) model's chain.
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
inst.load_lora(m, None, "a.safetensors", 1.0, 1.0)
first = m.get_attachment(self.ATTACH)
inst.load_lora(m, None, "b.safetensors", 1.0, 1.0)
second = m.get_attachment(self.ATTACH)
self.assertIsNot(first, second) # brand-new list
self.assertEqual(len(first), 1) # the earlier list is untouched
self.assertEqual(len(second), 2)
def test_double_install_is_idempotent(self):
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
lora_optimizer._install_lora_name_stamp(L) # second call is a no-op
inst = L()
m = _AttachHost()
inst.load_lora(m, None, "a.safetensors", 1.0, 1.0)
# Only one wrapper -> one entry per call (double-wrap would give two).
self.assertEqual(len(m.get_attachment(self.ATTACH)), 1)
def test_stamping_exception_is_swallowed(self):
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
class _Bad:
def get_attachment(self, key):
return None
def set_attachments(self, key, val):
raise RuntimeError("boom")
bad = _Bad()
out = inst.load_lora(bad, None, "a.safetensors", 1.0, 1.0)
# Original result still returned; loading is never broken by stamping.
self.assertEqual(out, (bad, None))
def test_clip_none_handled(self):
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
out = inst.load_lora(m, None, "a.safetensors", 1.0, 0.0)
self.assertIsNone(out[1])
self.assertEqual(len(m.get_attachment(self.ATTACH)), 1)
def test_clip_stamped_via_patcher(self):
# clip_out is a CLIP wrapper -> stamp lands on clip.patcher, which is
# exactly where the read path (clip.patcher.get_attachment) looks.
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
c = _ClipWrap()
inst.load_lora(m, c, "a.safetensors", 0.8, 0.6)
cs = c.patcher.get_attachment(self.ATTACH)
self.assertEqual(cs[0]["name"], "a.safetensors")
self.assertEqual(cs[0]["strength_clip"], 0.6)
def test_falsy_lora_name_not_stamped(self):
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
inst.load_lora(m, None, "", 1.0, 1.0)
self.assertIsNone(m.get_attachment(self.ATTACH))
def test_returns_original_result_object(self):
# The wrapper must return exactly what the original returned (same
# model/clip objects) so downstream nodes see the real patched models.
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
c = _ClipWrap()
out = inst.load_lora(m, c, "a.safetensors", 1.0, 1.0)
self.assertIs(out[0], m)
self.assertIs(out[1], c)
class _AttachModel(_FakePatcher):
"""ModelPatcher stand-in that also carries a loraopt_chain_names
attachment and copies it (shallow) on clone, like the real one."""
def __init__(self, patches=None, stamps=None):
super().__init__(patches)
self._att = {}
if stamps is not None:
self._att[lora_optimizer.LORAOPT_CHAIN_NAMES_ATTACH] = stamps
self.model = types.SimpleNamespace()
def get_attachment(self, key):
return self._att.get(key, None)
def set_attachments(self, key, val):
self._att[key] = val
def clone(self):
n = _AttachModel({k: v[:] for k, v in self.patches.items()})
n._att = dict(self._att)
return n
class _AttachClip:
"""CLIP stand-in whose stamps + patches live on its .patcher."""
def __init__(self, patcher):
self.patcher = patcher
self.cond_stage_model = types.SimpleNamespace()
def clone(self):
return _AttachClip(self.patcher.clone())
class TestInlineStampMatcher(unittest.TestCase):
"""Commit 2: _resolve_stamp_names aligns stamped filenames to chain groups
by ORDER and STRENGTH (per branch)."""
@staticmethod
def _m(groups, stamps, field):
return lora_optimizer.LoRAOptimizerInline._resolve_stamp_names(
groups, stamps, field)
def test_matches_two_groups_by_strength_model(self):
groups = [{"strength": 0.8, "entries": {}}, {"strength": 0.5, "entries": {}}]
stamps = [{"name": "a", "strength_model": 0.8, "strength_clip": 0.3},
{"name": "b", "strength_model": 0.5, "strength_clip": 0.4}]
self.assertEqual(self._m(groups, stamps, "strength_model"), ["a", "b"])
def test_matches_by_clip_strength_field(self):
groups = [{"strength": 0.6, "entries": {}}]
stamps = [{"name": "te", "strength_model": 0.0, "strength_clip": 0.6}]
self.assertEqual(self._m(groups, stamps, "strength_clip"), ["te"])
def test_strength_mismatch_leaves_group_unnamed(self):
groups = [{"strength": 0.9, "entries": {}}]
stamps = [{"name": "a", "strength_model": 0.2, "strength_clip": 1.0}]
self.assertEqual(self._m(groups, stamps, "strength_model"), [None])
def test_fewer_stamps_than_groups_degrades(self):
groups = [{"strength": 0.8, "entries": {}}, {"strength": 0.5, "entries": {}}]
stamps = [{"name": "a", "strength_model": 0.8, "strength_clip": 1.0}]
self.assertEqual(self._m(groups, stamps, "strength_model"), ["a", None])
def test_empty_stamps_all_unnamed(self):
groups = [{"strength": 0.8, "entries": {}}]
self.assertEqual(self._m(groups, [], "strength_model"), [None])
def test_malformed_stamp_entry_skipped(self):
groups = [{"strength": 0.8, "entries": {}}]
self.assertEqual(self._m(groups, ["not a dict"], "strength_model"), [None])
class TestInlineStampReadPath(unittest.TestCase):
"""Commit 2: execute_inline consumes stamped filenames -> real names in the
fingerprint + file identity on the virtual items (reconciling with the
file-based Stack community/memory dataset). Unstamped -> captured fallback."""
def _node(self):
return lora_optimizer.LoRAOptimizerInline()
@staticmethod
def _out(result):
return result["result"] if isinstance(result, dict) else result
def _capture_merge(self, node, seen):
def fake_merge(m, stack, output_strength, **kw):
seen["model"] = m
seen["stack"] = stack
seen.update(kw)
return (m, kw.get("clip"), "engine report", None, None)
node.optimize_merge = fake_merge
@staticmethod
def _resolves(kind, name):
return "/loras/" + name
def test_two_stamps_named_and_get_file_identity(self):
patches = _chain_patches((0.8, {"a": _adapter()}), (0.5, {"b": _adapter()}))
stamps = [{"name": "styles/a.safetensors", "strength_model": 0.8,
"strength_clip": 0.7},
{"name": "chars/b.safetensors", "strength_model": 0.5,
"strength_clip": 0.5}]
model = _AttachModel(patches, stamps)
node = self._node()
seen = {}
self._capture_merge(node, seen)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
side_effect=self._resolves):
result = node.execute_inline(model, output_strength=1.0)
report = self._out(result)[2]
# Real names surface in the fingerprint.
self.assertIn("styles/a.safetensors", report)
self.assertIn("chars/b.safetensors", report)
stack = seen["stack"]
self.assertEqual(stack[0]["_resolved_file_name"], "styles/a.safetensors")
self.assertEqual(stack[1]["_resolved_file_name"], "chars/b.safetensors")
# Salted display/merge-cache name is kept SEPARATE from the identity.
self.assertTrue(stack[0]["name"].startswith("chain lora #1"))
# Persistent identity is the FILE key (reconciles with Stack), not captured:.
k0 = lora_optimizer.LoRAAutoTuner._persistent_lora_key(stack[0])
self.assertEqual(k0, "styles/a.safetensors")
self.assertFalse(k0.startswith("captured:"))
def test_fewer_stamps_than_groups_leaves_extra_unnamed(self):
patches = _chain_patches((0.8, {"a": _adapter()}), (0.5, {"b": _adapter()}))
stamps = [{"name": "only/first.safetensors", "strength_model": 0.8,
"strength_clip": 1.0}]
model = _AttachModel(patches, stamps)
node = self._node()
seen = {}
self._capture_merge(node, seen)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
side_effect=self._resolves):
node.execute_inline(model, output_strength=1.0)
stack = seen["stack"]
self.assertEqual(stack[0]["_resolved_file_name"], "only/first.safetensors")
self.assertNotIn("_resolved_file_name", stack[1]) # unmatched -> unnamed
def test_strength_mismatch_group_stays_unnamed(self):
patches = _chain_patches((0.8, {"a": _adapter()}))
stamps = [{"name": "wrong.safetensors", "strength_model": 0.2,
"strength_clip": 1.0}]
model = _AttachModel(patches, stamps)
node = self._node()
seen = {}
self._capture_merge(node, seen)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
side_effect=self._resolves):
result = node.execute_inline(model, output_strength=1.0)
report = self._out(result)[2]
self.assertNotIn("wrong.safetensors", report)
self.assertNotIn("_resolved_file_name", seen["stack"][0])
def test_clip_only_leftover_group_gets_clip_name(self):
# TE-only LoRA: no model group, one clip group named by its clip stamp.
model = _AttachModel({}, [])
clip = _AttachClip(_AttachModel(
_chain_patches((0.6, {"te.a": _adapter()})),
[{"name": "te/style.safetensors", "strength_model": 0.0,
"strength_clip": 0.6}]))
node = self._node()
seen = {}
self._capture_merge(node, seen)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
side_effect=self._resolves):
result = node.execute_inline(model, output_strength=1.0, clip=clip)
report = self._out(result)[2]
self.assertIn("te/style.safetensors", report)
self.assertEqual(seen["stack"][0]["_resolved_file_name"],
"te/style.safetensors")
def test_no_attachment_keeps_generic_names_and_captured_identity(self):
# Regression: unstamped loader (fake patcher lacks the attachment API)
# -> generic "chain lora #N" name + captured: content-hash identity.
patches = _chain_patches((0.8, {"a": _adapter()}))
model = _FakePatcher(patches)
model.model = types.SimpleNamespace()
node = self._node()
seen = {}
self._capture_merge(node, seen)
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
result = node.execute_inline(model, output_strength=1.0)
report = self._out(result)[2]
self.assertIn("#1: 1 keys", report) # generic fingerprint line
stack = seen["stack"]
self.assertNotIn("_resolved_file_name", stack[0])
self.assertTrue(stack[0]["name"].startswith("chain lora #1"))
# captured identity requires get_full_path -> None (in-memory hash)
key = lora_optimizer.LoRAAutoTuner._persistent_lora_key(stack[0])
self.assertTrue(key.startswith("captured:"))
class TestInlineNamedReconciliation(unittest.TestCase):
"""Commit 2: a NAMED inline item shares the memory/community identity with
a Stack run of the same file (file-bytes content hash + name key)."""
def test_persistent_key_equals_stack_key(self):
real = "styles/a.safetensors"
named = {"name": "chain lora #1 [zzz]", "_precomputed_diffs": True,
"_resolved_file_name": real, "strength": 0.8, "lora": {}}
stack_item = {"name": real, "strength": 0.8, "lora": {}}
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
side_effect=lambda kind, name: "/loras/" + name):
k_named = lora_optimizer.LoRAAutoTuner._persistent_lora_key(named)
k_stack = lora_optimizer.LoRAAutoTuner._persistent_lora_key(stack_item)
self.assertEqual(k_named, k_stack)
self.assertEqual(k_named, real)
def test_content_hash_equals_stack_file_hash(self):
import tempfile
with tempfile.TemporaryDirectory() as memdir:
path = os.path.join(memdir, "a.safetensors")
with open(path, "wb") as f:
f.write(b"fake lora file bytes " * 64)
real = "styles/a.safetensors"
named = {"name": "chain lora #1 [q]", "_precomputed_diffs": True,
"_resolved_file_name": real, "strength": 0.8,
"lora": {"k": _adapter()}}
stack_item = {"name": real, "strength": 0.8, "lora": {}}
with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", memdir), \
mock.patch.object(lora_optimizer.folder_paths,
"get_full_path",
side_effect=lambda kind, name: path):
h_named = lora_optimizer.LoRAAutoTuner._lora_content_hash(named)
h_stack = lora_optimizer.LoRAAutoTuner._lora_content_hash(
stack_item)
# Named inline item hashes the FILE bytes (via _resolved_file_name),
# so it reconciles with the Stack path's file-bytes hash.
self.assertIsNotNone(h_named)
self.assertEqual(h_named, h_stack)
def test_unnamed_captured_still_hashes_in_memory(self):
# Regression: an item WITHOUT _resolved_file_name keeps the captured
# in-memory content hash (unchanged fallback).
item = {"name": "chain lora #1 [s]", "_precomputed_diffs": True,
"strength": 0.8, "lora": {"k": _adapter()}}
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
return_value=None):
h = lora_optimizer.LoRAAutoTuner._lora_content_hash(item)
key = lora_optimizer.LoRAAutoTuner._persistent_lora_key(item)
self.assertIsNotNone(h)
self.assertTrue(key.startswith("captured:"))
class TestInlineConservativeMatching(unittest.TestCase):
"""Important 2: _resolve_stamp_names must be conservative + order-
independent. A wrong _resolved_file_name pollutes the shared file-based
dataset, so name a group ONLY on an unambiguous unique-strength match."""
@staticmethod
def _m(groups, stamps, field):
return lora_optimizer.LoRAOptimizerInline._resolve_stamp_names(
groups, stamps, field)
@staticmethod
def _g(*strengths):
return [{"strength": s, "entries": {}} for s in strengths]
@staticmethod
def _s(*pairs):
return [{"name": n, "strength_model": sm, "strength_clip": sm}
for (n, sm) in pairs]
def test_same_strength_groups_both_unnamed(self):
# Two groups at the SAME strength -> ambiguous -> BOTH unnamed
# (better anonymous than a wrong file identity).
names = self._m(self._g(0.8, 0.8),
self._s(("a", 0.8), ("b", 0.8)), "strength_model")
self.assertEqual(names, [None, None])
def test_distinct_strengths_reversed_order_named_correctly(self):
# Reconstruction reversed the group order vs the stamp order: unique
# strengths must still map correctly by VALUE, not index.
names = self._m(self._g(0.5, 0.8),
self._s(("a", 0.8), ("b", 0.5)), "strength_model")
self.assertEqual(names, ["b", "a"])
def test_count_mismatch_unique_named_colliding_unnamed(self):
# An unstamped loader adds a 3rd group at 0.5: the unique 0.8 group is
# named; the two colliding 0.5 groups stay unnamed.
names = self._m(self._g(0.8, 0.5, 0.5),
self._s(("a", 0.8), ("b", 0.5)), "strength_model")
self.assertEqual(names, ["a", None, None])
def test_duplicate_stamp_strength_leaves_group_unnamed(self):
# Two stamps share a strength -> the matching group can't be resolved.
names = self._m(self._g(0.5),
self._s(("a", 0.5), ("b", 0.5)), "strength_model")
self.assertEqual(names, [None])
def test_near_equal_strengths_within_tol_treated_as_collision(self):
# 0.8 vs 0.80005 are within tol -> fuzzy-unique fails -> both unnamed
# (closes the 0.8/0.80001 false-positive window).
names = self._m(self._g(0.8, 0.80005),
self._s(("a", 0.8), ("b", 0.80005)), "strength_model")
self.assertEqual(names, [None, None])
def test_wrong_name_never_assigned_in_ambiguous_case(self):
names = self._m(self._g(0.8, 0.8),
self._s(("a", 0.8), ("b", 0.8)), "strength_model")
self.assertNotIn("a", names)
self.assertNotIn("b", names)
class TestInlineAmbiguousNoPollution(unittest.TestCase):
"""Integration: an ambiguous (same-strength) captured chain must get NO
file identity -> falls back to captured identity, so it can never write a
wrong file's stats into the shared dataset."""
def _node(self):
return lora_optimizer.LoRAOptimizerInline()
@staticmethod
def _out(result):
return result["result"] if isinstance(result, dict) else result
def test_same_strength_chain_gets_no_resolved_file_name(self):
patches = _chain_patches((0.8, {"a": _adapter()}),
(0.8, {"b": _adapter()}))
stamps = [{"name": "styles/a.safetensors", "strength_model": 0.8,
"strength_clip": 0.8},
{"name": "styles/b.safetensors", "strength_model": 0.8,
"strength_clip": 0.8}]
model = _AttachModel(patches, stamps)
node = self._node()
seen = {}
def fake_merge(m, stack, output_strength, **kw):
seen["stack"] = stack
return (m, kw.get("clip"), "engine report", None, None)
node.optimize_merge = fake_merge
with mock.patch.object(lora_optimizer.folder_paths, "get_full_path",
side_effect=lambda kind, name: "/loras/" + name):
result = node.execute_inline(model, output_strength=1.0)
report = self._out(result)[2]
# No real name displayed, no file identity attached.
self.assertNotIn("styles/a.safetensors", report)
self.assertNotIn("styles/b.safetensors", report)
for item in seen["stack"]:
self.assertNotIn("_resolved_file_name", item)
class TestLoraNameStampHardening(unittest.TestCase):
"""Review hardening: 0/0 early-return guard, signature-agnostic wrapper,
functools.wraps + __wrapped__-chain idempotency."""
ATTACH = "loraopt_chain_names"
def _fresh_loader_cls(self):
class _L:
def load_lora(self, model, clip, lora_name,
strength_model, strength_clip):
return (model, clip)
return _L
# --- Important 1: strength 0/0 loads must NOT stamp / mutate input ---
def test_zero_zero_load_adds_no_stamp(self):
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
inst.load_lora(m, None, "a.safetensors", 0.0, 0.0)
self.assertIsNone(m.get_attachment(self.ATTACH)) # no phantom stamp
def test_zero_zero_load_does_not_mutate_input_attachment(self):
# Stock LoraLoader early-returns the INPUT (model, clip) UNCHANGED on a
# 0/0 load (no clone) -> stamping would mutate a shared upstream model.
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
existing = [{"name": "up.safetensors", "strength_model": 1.0,
"strength_clip": 1.0}]
m.set_attachments(self.ATTACH, existing)
inst.load_lora(m, None, "a.safetensors", 0.0, 0.0)
after = m.get_attachment(self.ATTACH)
self.assertIs(after, existing) # exact same object, untouched
self.assertEqual(len(after), 1) # nothing appended
def test_nonzero_model_only_load_still_stamps(self):
# LoraLoaderModelOnly -> strength_clip=0 but strength_model!=0: NOT a
# 0/0 load, so it must still stamp.
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
inst.load_lora(m, None, "a.safetensors", 0.8, 0.0)
self.assertEqual(len(m.get_attachment(self.ATTACH)), 1)
# --- Important 3: signature-agnostic wrapper ---
def test_stamps_when_called_with_keyword_args(self):
# ComfyUI invokes node FUNCTION with keyword args from the inputs dict.
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
inst = L()
m = _AttachHost()
inst.load_lora(model=m, clip=None, lora_name="a.safetensors",
strength_model=0.8, strength_clip=0.6)
s = m.get_attachment(self.ATTACH)
self.assertEqual(s[0]["name"], "a.safetensors")
self.assertEqual(s[0]["strength_model"], 0.8)
self.assertEqual(s[0]["strength_clip"], 0.6)
def test_unknown_trailing_arg_does_not_break_loading_or_stamping(self):
# A future comfy adds a trailing param. The wrapper must pass it
# through (loading unbroken) AND still stamp from the leading args.
class _Future:
def load_lora(self, model, clip, lora_name, sm, sc, extra=None):
return (model, clip)
lora_optimizer._install_lora_name_stamp(_Future)
inst = _Future()
m = _AttachHost()
out = inst.load_lora(m, None, "a.safetensors", 0.8, 0.6, "new_param")
self.assertEqual(out, (m, None)) # loading unbroken
self.assertEqual(m.get_attachment(self.ATTACH)[0]["name"],
"a.safetensors")
def test_extraction_failure_skips_stamp_but_not_load(self):
# An exotic signature we can't extract from must never break loading.
class _Sig:
def load_lora(self, model, clip, lora_name):
return (model, clip)
lora_optimizer._install_lora_name_stamp(_Sig)
inst = _Sig()
m = _AttachHost()
out = inst.load_lora(m, None, "a.safetensors")
self.assertEqual(out, (m, None)) # no crash
self.assertIsNone(m.get_attachment(self.ATTACH)) # stamp skipped
# --- Minors: functools.wraps + __wrapped__-chain idempotency ---
def test_wrapper_preserves_metadata_via_functools_wraps(self):
class _Named:
def load_lora(self, model, clip, lora_name, sm, sc):
"orig docstring"
return (model, clip)
orig = _Named.load_lora
lora_optimizer._install_lora_name_stamp(_Named)
self.assertEqual(_Named.load_lora.__name__, "load_lora")
self.assertEqual(_Named.load_lora.__doc__, "orig docstring")
self.assertIs(_Named.load_lora.__wrapped__, orig)
def test_no_double_stamp_when_reinstalled_over_third_party_wrap(self):
L = self._fresh_loader_cls()
lora_optimizer._install_lora_name_stamp(L)
ours = L.load_lora
def third_party(self, *a, **k):
return ours(self, *a, **k)
third_party.__wrapped__ = ours # chain set, no _loraopt_stamped
L.load_lora = third_party
self.assertFalse(getattr(third_party, "_loraopt_stamped", False))
# Re-install (e.g. HotReload re-import) must find us in the chain and
# refuse to re-wrap.
self.assertTrue(lora_optimizer._install_lora_name_stamp(L))
self.assertIs(L.load_lora, third_party)
inst = L()
m = _AttachHost()
inst.load_lora(m, None, "a.safetensors", 1.0, 1.0)
self.assertEqual(len(m.get_attachment(self.ATTACH)), 1) # single stamp
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink