From 23961853936e8d73528410601f0d3884f73c9716 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 02:44:42 +0100 Subject: [PATCH 01/23] feat(lora): dense stack modes for multi-network sets Add lora_stack_mode with ties, dare_ties, dare_linear and magnitude_prune combination of per-network deltas when several loaded networks target one layer; sum stays the default and the exact factor path. Combined deltas ride the existing tail: hosted svd on sub-8-bit SDNQ, requantize at int8 and above, direct add elsewhere. Text-encoder layers and single-network sets keep plain summation. DARE masks draw from per-layer sha256 seeds so re-applies and cache entries stay deterministic; the stack settings join the activation and factor-cache signatures so settings changes re-apply without a reload. --- modules/lora/extra_networks_lora.py | 4 +- modules/lora/lora_apply.py | 17 +- modules/lora/lora_factor_cache.py | 2 + modules/lora/lora_sdnq.py | 37 ++-- modules/lora/lora_stack.py | 250 ++++++++++++++++++++++++++++ modules/lora/networks.py | 5 +- modules/ui_definitions.py | 6 + test/test-sdnq-lora-factors.py | 168 ++++++++++++++++++- ui/locale/locale_en.json | 4 + 9 files changed, 471 insertions(+), 22 deletions(-) create mode 100644 modules/lora/lora_stack.py diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index d0a6ffe85..f5cf5647d 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -178,8 +178,8 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): return [f'{name}:{te}:{unet}' for name, te, unet in zip(names, te_multipliers, unet_multipliers, strict=False)] def changed(self, requested: list[str], include: list[str] | None = None, exclude: list[str] | None = None) -> tuple[bool, str]: - from modules.lora import lora_sdnq - requested = requested + [f'stack={lora_sdnq.signature()}'] # settings-only mechanism changes must re-trigger activation + from modules.lora import lora_sdnq, lora_stack + requested = requested + [f'stack={lora_stack.signature()}{lora_sdnq.signature()}'] # settings-only stack or mechanism changes must re-trigger activation if shared.opts.lora_force_reload: debug_log(f'Network check: type=LoRA requested={requested} status="forced"') return True, "forced" diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 3cb8688c0..8367d61ce 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -5,6 +5,7 @@ import time from typing import TYPE_CHECKING import torch from modules.lora import lora_common as l +from modules.lora import lora_stack from modules import shared, devices, errors from modules.logger import log @@ -75,6 +76,9 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou pass batch_updown = None batch_ex_bias = None + stack_deltas = None + if lora_stack.mode() in lora_stack.DENSE_MODES and network_layer_name is not None and not network_layer_name.startswith('lora_te'): + stack_deltas = [] # collect per-net deltas; combined after the loop (bias deltas stay summed) loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks for net in loaded: module = net.modules.get(network_layer_name, None) @@ -107,7 +111,9 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou del weight if updown is not None: - if batch_updown is not None: + if stack_deltas is not None: + stack_deltas.append((net.name, updown.to(devices.device))) + elif batch_updown is not None: batch_updown += updown.to(batch_updown.device) else: batch_updown = updown.to(devices.device) @@ -136,6 +142,15 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou if elimit is not None: elimit() continue + if stack_deltas is not None and stack_deltas: + if len(stack_deltas) >= 2: + t0 = time.time() + batch_updown = lora_stack.combine(stack_deltas, network_layer_name) + l.timer.calc += time.time() - t0 + else: + batch_updown = stack_deltas[0][1] + if shared.opts.diffusers_offload_mode == "sequential": + batch_updown = batch_updown.to(devices.cpu) return batch_updown, batch_ex_bias diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index 2b00bbcc7..d6c95d681 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -51,10 +51,12 @@ def signature(wanted_names): if model_name is None: return None calib_path = lora_calib.calib_file(model_name) + from modules.lora import lora_stack parts = { 'model': model_name, 'rank': int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0), 'calib': int(os.path.getmtime(calib_path)) if lora_calib.enabled() and os.path.isfile(calib_path) else None, # the toggle is part of the identity: factors computed under the other setting must not replay + 'stack': lora_stack.signature(), 'nets': [], } for name, te, unet, dyn in wanted_names: diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index d95849527..7973121b5 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -47,7 +47,7 @@ a low-rank delta hosts exactly however fat it is. import torch from modules import devices, shared -from modules.lora import lora_calib, lora_factor_cache +from modules.lora import lora_calib, lora_factor_cache, lora_stack from modules.lora import lora_common as l from modules.logger import log @@ -137,6 +137,9 @@ def factor_candidate(self, network_layer_name, wanted_names): return False # declined layers with factors still attached are stripped by the activate fallthrough if getattr(self, 'sdnq_dequantizer', None) is None or self.__class__.__name__ != 'SDNQLinear': return False + if wanted_names != () and lora_stack.mode() in lora_stack.DENSE_MODES and not network_layer_name.startswith('lora_te'): + if sum(1 for net in l.loaded_networks if net.modules.get(network_layer_name, None) is not None) >= 2: + return False # dense stack modes combine dense deltas; the factor concat would sum if hasattr(self, 'sdnq_lora_svd_stash'): return True if wanted_names == (): # nothing attached, nothing to remove @@ -277,14 +280,16 @@ def apply_cached(self, network_layer_name, wanted_names): deq = self.sdnq_dequantizer dtype = deq.result_dtype remove_factors(self) # before the rule: the svd-channel check must see the checkpoint's own state, and a declined layer must fall through pristine + stack_dense = lora_stack.mode() in lora_stack.DENSE_MODES and not network_layer_name.startswith('lora_te') members = [] - for net in l.loaded_networks: - module = net.modules.get(network_layer_name, None) - if module is None: - continue - factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) - if factors is not None: - members.append(factors) + if not stack_dense: + for net in l.loaded_networks: + module = net.modules.get(network_layer_name, None) + if module is None: + continue + factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) + if factors is not None: + members.append(factors) if len(members) == 0 and self.svd_up is None: step = float(self.scale.detach().float().mean()) if step > 0 and rms / step > REQUANT_RATIO and energy < REQUANT_ENERGY: @@ -328,13 +333,15 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): dtype = deq.result_dtype members = [] - for net in l.loaded_networks: - module = net.modules.get(network_layer_name, None) - if module is None: - continue - factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) - if factors is not None: - members.append(factors) + stack_dense = lora_stack.mode() in lora_stack.DENSE_MODES and not network_layer_name.startswith('lora_te') + if not stack_dense: # dense stack modes host the combined delta wholesale; the members' content is already inside it + for net in l.loaded_networks: + module = net.modules.get(network_layer_name, None) + if module is None: + continue + factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) + if factors is not None: + members.append(factors) # requantize keeps a delta the grid can resolve and that truncation would genuinely # cut: both terms must agree, since a thin delta rounds away on the grid however diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py new file mode 100644 index 000000000..224505cab --- /dev/null +++ b/modules/lora/lora_stack.py @@ -0,0 +1,250 @@ +"""Stack modes for combining multiple LoRA networks beyond plain summation. + +Dense modes (ties, dare_ties, dare_linear, magnitude_prune) combine the +networks' dense deltas elementwise; the result rides the normal apply tail +(side-channel hosting on sub-8-bit SDNQ, requantize at int8 and above, +direct add on unquantized layers). Select modes (klora, estlora) keep both +networks' contributions separate and choose a per-layer winner, shifting +from the first loaded network (subject) toward the second (style) across +the sampling steps. Selection scores depend only on the weights, so the +shift reduces to at most one flip per layer per generation, executed from +the step callback against a schedule finalized at apply time. + +TIES arXiv:2306.01708, DARE arXiv:2311.03099, K-LoRA arXiv:2502.18461, +EST-LoRA arXiv:2508.02165 (its measured style-discrepancy estimate is +exposed as an option instead of being derived from probe generations). +""" + +import weakref +import hashlib + +import torch + +from modules import shared +from modules.logger import log + + +DENSE_MODES = ('ties', 'dare_ties', 'dare_linear', 'magnitude_prune') +SELECT_MODES = ('klora', 'estlora') +KLORA_BETA = 0.5 # the paper's fixed ramp offset; only the slope is user-tunable +ROW_CHUNK = 512 # fp32 interiors run in first-dim slices; also fixes the DARE draw sequence +SAMPLE_CAP = 1 << 22 # strided subsample bound for magnitude quantiles (full-size quantile exceeds torch limits) + +state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'total_steps': 0, 'finalized': False} +warned: set = set() + + +def mode(): + return getattr(shared.opts, 'lora_stack_mode', 'sum') or 'sum' + + +def density(): + return float(getattr(shared.opts, 'lora_stack_density', 0.5)) + + +def ramp_alpha(): + return float(getattr(shared.opts, 'lora_stack_alpha', 1.5)) + + +def manual_discrepancy(): + return float(getattr(shared.opts, 'lora_stack_discrepancy', 0.5)) + + +def signature(): + m = mode() + if m in DENSE_MODES: + return f'{m}:{density():.2f}' + if m in SELECT_MODES: + return f'{m}:{ramp_alpha():.2f}:{manual_discrepancy():.2f}' + return 'sum' + + +def warn_once(key, message): + if key not in warned: + warned.add(key) + log.warning(message) + + +def select_blocked(): + return 'Model' in (getattr(shared.opts, 'cuda_compile', None) or []) + + +def active_dense(n_contrib): + return mode() in DENSE_MODES and n_contrib >= 2 + + +def active_select(n_loaded): + m = mode() + if m not in SELECT_MODES: + return False + if n_loaded != 2: + warn_once('select-count', f'Network stack: mode={m} networks={n_loaded} requires exactly 2, using sum') + return False + if select_blocked(): + warn_once('select-compile', f'Network stack: mode={m} disabled with model compile, using sum') + return False + return True + + +def seed_for(layer_name, net_name): + payload = f'{layer_name}|{net_name}|{mode()}|{round(density(), 6)}' + return int.from_bytes(hashlib.sha256(payload.encode()).digest()[:8], 'little') + + +def magnitude_threshold(delta, dens): + flat = delta.abs().flatten() + step = max(1, flat.numel() // SAMPLE_CAP) + return torch.quantile(flat[::step].float(), 1.0 - dens) + + +def dare_generator(device, layer_name, net_name): + gen = torch.Generator(device=device) + gen.manual_seed(seed_for(layer_name, net_name)) + return gen + + +def combine(named_deltas, layer_name): + """Combine per-network dense deltas under the active dense mode; returns a tensor in the first delta's dtype.""" + m = mode() + dens = density() + deltas = [d for _, d in named_deltas] + out_dtype = deltas[0].dtype + result = torch.zeros_like(deltas[0], dtype=torch.float32) + thresholds = [magnitude_threshold(d, dens) for d in deltas] if m in ('ties', 'magnitude_prune') else [None] * len(deltas) + gens = [dare_generator(deltas[0].device, layer_name, name) for name, _ in named_deltas] if m in ('dare_ties', 'dare_linear') else [None] * len(deltas) + for start in range(0, deltas[0].shape[0], ROW_CHUNK): + stop = min(start + ROW_CHUNK, deltas[0].shape[0]) + chunks = [] + for i, d in enumerate(deltas): + c = d[start:stop].to(torch.float32) + if thresholds[i] is not None: + c = c * (c.abs() >= thresholds[i]) + if gens[i] is not None: + keep = torch.rand(c.shape, generator=gens[i], device=c.device, dtype=torch.float32) < dens + c = c * keep / dens + chunks.append(c) + if m in ('ties', 'dare_ties'): + total = torch.stack(chunks).sum(dim=0) + elected = torch.sign(total) + agree = [c * ((torch.sign(c) == elected) & (c != 0)) for c in chunks] + count = torch.stack([(a != 0).to(torch.float32) for a in agree]).sum(dim=0).clamp(min=1.0) + result[start:stop] = torch.stack(agree).sum(dim=0) / count + else: # dare_linear, magnitude_prune: independent per-delta edits, plain sum + result[start:stop] = torch.stack(chunks).sum(dim=0) + return result.to(out_dtype) + + +def score_topk(up, down, k): + """K-LoRA layer score: sum of the top-K absolute delta entries (one dense materialization).""" + d = (up.to(torch.float32) @ down.to(torch.float32)).abs().flatten() + values = torch.topk(d, min(int(k), d.numel()), sorted=False).values + return float(values.sum()), float(d.sum()) + + +def score_energy(up, down): + """EST layer score: squared Frobenius norm of up@down via the Gram identity, no materialization.""" + u = up.to(torch.float32) + dn = down.to(torch.float32) + return float(((u.t() @ u) * (dn @ dn.t())).sum()) + + +def clear(): + state['entries'] = {} + state['flips'] = {} + state['gamma'] = 1.0 + state['total_steps'] = 0 + state['finalized'] = False + + +def register(layer_name, module, kind, segments, scores, factors=None): + """Record a select-mode layer: its two segments (or bf16 factor pairs) and static scores. + + kind 'factor': segments = [(start, stop), (start, stop)] column ranges in svd_up/svd_down + with the transposed-layout flag appended; stashes both segments' values for flips. + kind 'weight': factors = [(up0, down0), (up1, down1)] kept for recompute-from-backup. + """ + entry = {'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'factors': factors, 'stash': None} + if kind == 'factor': + (s0, s1), (t0, t1), transposed = segments + up = module.svd_up.data + entry['stash'] = (segment_view(up, s0, s1, transposed).clone(), segment_view(up, t0, t1, transposed).clone()) + state['entries'][layer_name] = entry + state['finalized'] = False + + +def segment_view(up, start, stop, transposed): + return up[start:stop] if transposed else up[:, start:stop] + + +def layer_flip_step(scores, total_steps): + """First step index at which the style side wins; total_steps when it never does, 0 when style wins from the start.""" + m = mode() + sc, ss = scores + for step in range(total_steps): + t = step / max(1, total_steps - 1) + if m == 'klora': + ramp = state['gamma'] * (ramp_alpha() * t + KLORA_BETA) + if ss * ramp > sc: + return step + else: # estlora: content keeps the layer while sc >= gamma_t * ss + ramp = ramp_alpha() * t + (1.0 - manual_discrepancy()) + if sc < ramp * ss: + return step + return total_steps + + +def finalize(total_steps): + """Build the inverted flip map for the pass; select-mode layers start at their step-0 winner.""" + state['total_steps'] = int(total_steps) + state['flips'] = {} + for layer_name, entry in state['entries'].items(): + flip_at = layer_flip_step(entry['scores'], state['total_steps']) + initial = 1 if flip_at == 0 else 0 + apply_selection(layer_name, entry, initial) + if 0 < flip_at < state['total_steps']: + state['flips'].setdefault(flip_at, []).append(layer_name) + state['finalized'] = True + + +def reset(total_steps): + """Per-pass reset from set_callbacks_p: restore initial selections and reschedule for this pass's step count.""" + if mode() not in SELECT_MODES or not state['entries']: + return + finalize(total_steps) + + +def on_step(step): + """Flip the layers whose crossover is this step; non-flip steps are a dict miss.""" + if not state['finalized']: + return + for layer_name in state['flips'].get(int(step), ()): + entry = state['entries'].get(layer_name) + if entry is not None: + apply_selection(layer_name, entry, 1) + + +def apply_selection(layer_name, entry, winner): + module = entry['module']() + if module is None: + state['entries'].pop(layer_name, None) + return + if entry['kind'] == 'factor': + (s0, s1), (t0, t1), transposed = entry['segments'] + up = module.svd_up.data + keep, drop = ((t0, t1), (s0, s1)) if winner == 1 else ((s0, s1), (t0, t1)) + stash = entry['stash'][winner] + segment_view(up, keep[0], keep[1], transposed).copy_(stash.to(device=up.device, dtype=up.dtype)) + segment_view(up, drop[0], drop[1], transposed).zero_() + else: + weight_selection(module, entry, winner) + + +def weight_selection(module, entry, winner): + backup = getattr(module, 'network_weights_backup', None) + if not isinstance(backup, torch.Tensor): # fuse mode keeps a bool sentinel, not a pristine copy + warn_once('select-nobackup', 'Network stack: select flip skipped, no weight backup') + return + up, down = entry['factors'][winner] + weight = backup.to(device=module.weight.device, dtype=torch.float32) + delta = up.to(device=module.weight.device, dtype=torch.float32) @ down.to(device=module.weight.device, dtype=torch.float32) + module.weight.data.copy_((weight + delta.reshape(weight.shape)).to(module.weight.dtype)) diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 9e3b2e7a0..58548b82a 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -5,6 +5,7 @@ from modules.errorlimiter import limit_errors from modules.lora import lora_common as l from modules.lora import lora_overrides from modules.lora import lora_sdnq +from modules.lora import lora_stack from modules.lora.lora_apply import network_apply_weights, network_apply_direct, network_backup_weights, network_calc_weights from modules import shared, devices, sd_models from modules.logger import log, console @@ -89,7 +90,7 @@ def network_activate(include=None, exclude=None): refused = 0 with devices.inference_context(), pbar: wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () - stack_sig = lora_sdnq.signature() # apply-mechanism token tracked beside network_current_names so a settings-only flip re-applies + stack_sig = lora_stack.signature() + lora_sdnq.signature() # tracked beside network_current_names so settings-only stack or mechanism changes re-apply applied_layers.clear() lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind lora_sdnq.hosted_layers.clear() @@ -100,7 +101,7 @@ def network_activate(include=None, exclude=None): for _, module in modules[component]: network_layer_name = getattr(module, 'network_layer_name', None) current_names = getattr(module, "network_current_names", ()) - if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == component_wanted and getattr(module, 'network_current_stack', '') == stack_sig): + if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == component_wanted and getattr(module, 'network_current_stack', 'sum') == stack_sig): if task is not None: pbar.update(task, advance=1) continue diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index a873ea68d..93ebbb3aa 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -709,6 +709,12 @@ def create_settings(cmd_opts): "lora_sdnq_host_calib": OptionInfo(True, "LoRA quantized host calibration"), "lora_sdnq_host_cache": OptionInfo(10, "LoRA quantized host cache", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}), + "lora_stack_sep": OptionInfo("

Stacking options

", "", gr.HTML), + "lora_stack_mode": OptionInfo("sum", "LoRA stack mode", gr.Dropdown, {"choices": ["sum", "ties", "dare_ties", "dare_linear", "magnitude_prune", "klora", "estlora"]}), + "lora_stack_density": OptionInfo(0.5, "LoRA stack density", gr.Slider, {"minimum": 0.05, "maximum": 1.0, "step": 0.05}), + "lora_stack_alpha": OptionInfo(1.5, "LoRA stack ramp", gr.Slider, {"minimum": 0.0, "maximum": 3.0, "step": 0.1}), + "lora_stack_discrepancy": OptionInfo(0.5, "LoRA stack discrepancy", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}), + "lora_meta_sep": OptionInfo("

Metadata

", "", gr.HTML), "lora_add_hashes_to_infotext": OptionInfo(False, "LoRA add hash info to metadata"), })) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 3b745ad11..1949fd6a3 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -76,7 +76,7 @@ modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) from modules.errors import log # pylint: disable=wrong-import-position from modules import shared, sd_models # pylint: disable=wrong-import-position -from modules.lora import network, network_lora, lora_sdnq, networks # pylint: disable=wrong-import-position +from modules.lora import network, network_lora, lora_sdnq, lora_stack, networks # pylint: disable=wrong-import-position from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position from sdnq.quantizer import sdnq_quantize_layer, SDNQConfig # pylint: disable=wrong-import-position @@ -1403,6 +1403,165 @@ def test_cache_fastpath_serves_mixed_set(): return True +def test_cache_fastpath_serves_dense_pair(): + import tempfile + layer = build_layer('uint4') + with tempfile.TemporaryDirectory() as tmp: + with host_rank(64), host_cache(10, os.path.join(tmp, 'cache')), stack_mode('ties'), mock_model(lin=layer): + net1, _D1 = cache_fixture(tmp, layer, name='densea', seed=65) + net2, _D2 = cache_fixture(tmp, layer, name='denseb', seed=66) + activate(net1, net2) + first = dq(layer) + activate() + with counting_calc() as calls: + activate(net1, net2) # the dense combine lives inside delta assembly; the replay skips both + assert calls['n'] == 0, 'a dense-pair replay must not assemble or combine deltas' + assert hasattr(layer, 'sdnq_lora_svd_stash') + assert torch.equal(dq(layer), first), 'dense-pair replay must be bit-identical to the fresh apply' + activate() + return True + + +CAT_STACK = category('stack-dense') + + +@contextmanager +def stack_mode(name, dens=None): + old_m = getattr(shared.opts, 'lora_stack_mode', 'sum') + old_d = getattr(shared.opts, 'lora_stack_density', 0.5) + shared.opts.lora_stack_mode = name + if dens is not None: + shared.opts.lora_stack_density = dens + try: + yield + finally: + shared.opts.lora_stack_mode = old_m + shared.opts.lora_stack_density = old_d + + +def test_ties_sign_consensus_drops_conflicts(): + with stack_mode('ties', dens=1.0): # density 1 disables the trim, isolating sign election + d1 = torch.tensor([[1.0, 1.0, -1.0]], device=DEVICE) + d2 = torch.tensor([[2.0, -0.5, -2.0]], device=DEVICE) + out = lora_stack.combine([('a', d1), ('b', d2)], 'lora_transformer_test') + expected = torch.tensor([[1.5, 1.0, -1.5]], device=DEVICE) # agree: mean; conflict: majority-mass side only + assert torch.allclose(out, expected), f'{out.tolist()}' + return True + + +def test_dare_mask_is_deterministic_across_calls(): + torch.manual_seed(21) + d1 = torch.randn(64, 96, device=DEVICE) * 1e-2 + d2 = torch.randn(64, 96, device=DEVICE) * 1e-2 + with stack_mode('dare_linear', dens=0.5): + out1 = lora_stack.combine([('a', d1), ('b', d2)], 'lora_transformer_test') + out2 = lora_stack.combine([('a', d1), ('b', d2)], 'lora_transformer_test') + other = lora_stack.combine([('a', d1), ('b', d2)], 'lora_transformer_other') + assert torch.equal(out1, out2), 'same layer and nets must draw the same masks' + assert not torch.equal(out1, other), 'a different layer must draw different masks' + return True + + +def test_dare_rescales_by_inverse_density(): + torch.manual_seed(22) + d1 = torch.randn(64, 96, device=DEVICE) + d2 = torch.randn(64, 96, device=DEVICE) + with stack_mode('dare_linear', dens=0.5): + out = lora_stack.combine([('a', d1), ('b', d2)], 'lora_transformer_test') + cands = torch.stack([torch.zeros_like(d1), 2 * d1, 2 * d2, 2 * d1 + 2 * d2]) + nearest = (cands - out.unsqueeze(0)).abs().min(dim=0).values + assert float(nearest.max()) < 1e-5, 'every element must be a 1/density-rescaled subset sum' + zero_frac = float((out == 0).float().mean()) + assert 0.1 < zero_frac < 0.45, f'both-dropped fraction {zero_frac} should sit near 0.25' + return True + + +def test_magnitude_prune_keeps_top_density(): + torch.manual_seed(23) + d1 = torch.randn(128, 64, device=DEVICE) + d2 = torch.zeros_like(d1) # inert second delta isolates the trim + with stack_mode('magnitude_prune', dens=0.25): + out = lora_stack.combine([('a', d1), ('b', d2)], 'lora_transformer_test') + kept = out != 0 + frac = float(kept.float().mean()) + assert 0.2 < frac < 0.3, f'kept fraction {frac}' + assert torch.equal(out[kept], d1[kept]), 'kept elements must pass through unchanged' + assert float(d1.abs()[~kept].max()) <= float(d1.abs()[kept].min()) + 1e-6, 'kept set must be the top magnitudes' + return True + + +def test_dense_two_plain_loras_hosted_not_summed(): + layer = build_layer('uint4') + A1, B1, D1 = make_delta(seed=31, sigma=1e-2) + A2, B2, D2 = make_delta(seed=32, sigma=1e-2) + n1 = make_net('td1', layer, A1, B1) + n2 = make_net('td2', layer, A2, B2) + with host_rank(64), mock_model(lin=layer): + Wdq0 = dq(layer) + with stack_mode('ties', dens=0.5): + activate(n1, n2) + # hosted at rank 64 leaves a rank-64 factor bucket; the exact concat of two rank-8 nets would leave 16 + assert layer.svd_up.shape[1] == 64, f'dense mode must route a factorable pair to hosting, rank={layer.svd_up.shape[1]}' + eff = dq(layer) - Wdq0 + activate() + assert torch.equal(dq(layer), Wdq0), 'removal must restore bit-exact' + with stack_mode('ties', dens=0.5): + ref = lora_stack.combine([('td1', D1), ('td2', D2)], 'lora_transformer_test') + s = D1 + D2 + assert float((eff - s).norm() / s.norm()) > 0.05, 'ties result must differ from the plain sum' + assert rho_of(eff, ref) > 0.8, f'hosted ties delta must track the ties reference, rho={rho_of(eff, ref):.3f}' # rank-64 truncation of the densified delta keeps ~0.89 + assert float((eff - ref).norm()) < float((eff - s).norm()), 'hosted result must sit closer to the ties reference than to the plain sum' + return True + + +def test_single_net_ignores_dense_mode(): + layer = build_layer('uint4') + A, B, D = make_delta(seed=33) + net = make_net('solo', layer, A, B) + with mock_model(lin=layer), stack_mode('ties', dens=0.5): + Wdq0 = dq(layer) + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'single net must stay on the exact factor path' + assert rho_of(dq(layer) - Wdq0, D) > 0.99 + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_te_layer_stays_plain_sum(): + layer = build_layer('uint4') + layer.network_layer_name = 'lora_te_test' + A1, B1, D1 = make_delta(seed=34) + A2, B2, D2 = make_delta(seed=35) + n1 = make_net('te1', layer, A1, B1) + n2 = make_net('te2', layer, A2, B2) + with mock_model(lin=layer), stack_mode('ties', dens=0.5): + Wdq0 = dq(layer) + activate(n1, n2) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'te layers must stay on the exact concat path' + assert rho_of(dq(layer) - Wdq0, D1 + D2) > 0.99 + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_sum_mode_keeps_exact_stacking(): + layer = build_layer('uint4') + A1, B1, D1 = make_delta(seed=36) + A2, B2, D2 = make_delta(seed=37) + n1 = make_net('s1', layer, A1, B1) + n2 = make_net('s2', layer, A2, B2) + with mock_model(lin=layer), stack_mode('sum'): + Wdq0 = dq(layer) + activate(n1, n2) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'sum mode must keep the exact concat path' + assert layer.svd_up.shape[1] == 16, f'sum mode must concat exactly, rank={layer.svd_up.shape[1]}' + assert rho_of(dq(layer) - Wdq0, D1 + D2) > 0.99 + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + CAT_COMPILE = category('compile') @@ -1579,9 +1738,14 @@ def run_tests(): log.warning('=== Factor cache ===') for fn in [test_factor_cache_roundtrip_bitexact, test_factor_cache_invalidates_on_multiplier, test_factor_cache_int8_quantization, test_factor_cache_disabled_at_zero, test_factor_cache_invalidates_on_calib_toggle, - test_cache_fastpath_skips_calc, test_cache_fastpath_serves_mixed_set, + test_cache_fastpath_skips_calc, test_cache_fastpath_serves_mixed_set, test_cache_fastpath_serves_dense_pair, test_attach_trims_stored_null_tail]: run_test(CAT_FCACHE, fn) + log.warning('=== Stack modes: dense ===') + for fn in [test_ties_sign_consensus_drops_conflicts, test_dare_mask_is_deterministic_across_calls, test_dare_rescales_by_inverse_density, + test_magnitude_prune_keeps_top_density, test_dense_two_plain_loras_hosted_not_summed, test_single_net_ignores_dense_mode, + test_te_layer_stays_plain_sum, test_sum_mode_keeps_exact_stacking]: + run_test(CAT_STACK, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: run_test(CAT_COMPILE, fn) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 135283560..522ae19e6 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -868,6 +868,10 @@ {"id":"","label":"LoRA quantized host rank","localized":"","hint":"Maximum rank used to carry non-factorable adapter types (LoKR, LoHA, OFT, DoRA) on the side-channel of SDNQ models quantized below 8 bits, where merging would erase most of the adapter. Higher values keep more of the adapter at proportionally higher memory cost. Set to 0 to disable and merge into the quantized weights instead.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized host calibration","localized":"","hint":"Collection of per-channel activation statistics during normal generations on models quantized below 8 bits, cached per checkpoint. When available, the statistics weight side-channel hosting of non-factorable adapter types toward the channels carrying the most activation energy, improving delivered fidelity at the same host rank.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized host cache","localized":"","hint":"Disk budget in GB for caching the side-channel factors computed when hosting non-factorable adapter types on quantized models. A cached configuration skips the truncation math on reapply; least recently used entries are evicted past the budget. Set to 0 to disable.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack mode","localized":"","hint":"Combination rule applied when multiple LoRA networks target the same layer. sum adds all contributions. ties merges only elements whose signs agree after keeping the strongest ones; dare variants randomly drop elements and rescale the survivors; magnitude_prune keeps only the strongest elements of each network. klora and estlora select one of exactly two networks per layer, with the first network in the prompt treated as subject and the second as style, and the balance shifting from subject toward style over the sampling steps. Applied by the native loader only; other load methods behave as sum.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack density","localized":"","hint":"Fraction of elements kept by the ties, dare and magnitude_prune stack modes. Lower values keep only the strongest contributions of each network.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack ramp","localized":"","hint":"Steepness of the subject-to-style shift over the sampling steps for the klora and estlora stack modes. Higher values strengthen the late-step shift toward the style network.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack discrepancy","localized":"","hint":"Stand-in for the estlora mode's data-derived style separation estimate. Higher values shift the overall balance toward the subject network.","ui":"settings_lora"}, {"id":"","label":"LoRA auto-apply tags","localized":"","hint":"Automatically add trigger words/tags from LoRA metadata to your prompt.
Set to the number of tags to auto-apply, e.g., 3 = add top 3 trigger tags.
Set to 0 to disable, -1 to add all available tags.","ui":"settings_lora"}, {"id":"","label":"LoRA memory cache","localized":"","hint":"How many LoRAs to keep in network for future use before requiring reloading from storage","ui":"settings_lora"}, {"id":"","label":"LoRA add hash info to metadata","localized":"","hint":"Include LoRA file hashes in generated image metadata.
Useful for reproducibility and tracking which exact LoRA versions were used.","ui":"settings_lora"}, From 259e15fafe3f93ef61126f5ade3056b3e250b33e Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 02:58:11 +0100 Subject: [PATCH 02/23] feat(lora): per-layer select stack modes klora and estlora Two-network subject+style sets select a winner per layer instead of summing: scores are top-K magnitude sums (klora) or Frobenius energies (estlora), and a timestep ramp shifts layers from the subject network toward the style network across sampling, reduced to at most one precomputed flip per layer per pass. On sub-8-bit SDNQ the pair rides the side-channel as separate segments flipped in place; other layers recompute the winner from the pristine backup, so select modes force backup mode. Selection resets per pass from the callback setup and is gated off under model compile. estlora's measured style-discrepancy term is exposed as an option. Adds XYZ axes for the stack settings. --- modules/lora/extra_networks_lora.py | 3 + modules/lora/lora_apply.py | 8 +- modules/lora/lora_overrides.py | 3 + modules/lora/lora_sdnq.py | 67 +++++++- modules/lora/lora_stack.py | 81 ++++++++-- modules/lora/networks.py | 33 ++++ modules/processing_callbacks.py | 6 + scripts/xyz/xyz_grid_classes.py | 16 ++ test/test-sdnq-lora-factors.py | 228 ++++++++++++++++++++++++++++ 9 files changed, 426 insertions(+), 19 deletions(-) diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index f5cf5647d..c81b94072 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -225,6 +225,9 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): reason = '' load_method, load_reason = lora_overrides.get_method() + from modules.lora import lora_stack + if load_method != 'native' and lora_stack.mode() != 'sum': + lora_stack.warn_once(f'method-{load_method}', f'Network stack: mode={lora_stack.mode()} method={load_method} unsupported, using sum') if debug: import sys fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 8367d61ce..11b91d9e9 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -68,7 +68,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr return backup_size -def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, use_previous: bool = False, *, elimit: Callable[[], None] | None = None): +def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, use_previous: bool = False, *, elimit: Callable[[], None] | None = None, per_net: bool = False): if shared.opts.diffusers_offload_mode == "none": try: self.to(devices.device) @@ -77,8 +77,8 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou batch_updown = None batch_ex_bias = None stack_deltas = None - if lora_stack.mode() in lora_stack.DENSE_MODES and network_layer_name is not None and not network_layer_name.startswith('lora_te'): - stack_deltas = [] # collect per-net deltas; combined after the loop (bias deltas stay summed) + if per_net or (lora_stack.mode() in lora_stack.DENSE_MODES and network_layer_name is not None and not network_layer_name.startswith('lora_te')): + stack_deltas = [] # collect per-net deltas; combined after the loop unless the caller wants them separate (bias deltas stay summed) loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks for net in loaded: module = net.modules.get(network_layer_name, None) @@ -142,6 +142,8 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou if elimit is not None: elimit() continue + if per_net: + return stack_deltas, batch_ex_bias if stack_deltas is not None and stack_deltas: if len(stack_deltas) >= 2: t0 = time.time() diff --git a/modules/lora/lora_overrides.py b/modules/lora/lora_overrides.py index eb629896e..078b285b2 100644 --- a/modules/lora/lora_overrides.py +++ b/modules/lora/lora_overrides.py @@ -112,6 +112,9 @@ def disable_fuse(): round-trips it through its storage format. On quantized weights that is a dequantize-add-requantize cycle per network swap whose error compounds. """ + from modules.lora import lora_stack + if lora_stack.mode() in lora_stack.SELECT_MODES: + return True # select stack modes flip per-layer winners against the pristine backup sd_model = getattr(shared.sd_model, 'pipe', shared.sd_model) if is_quantized(sd_model): return True diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 7973121b5..b1ef10d45 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -174,6 +174,7 @@ def remove_factors(self): self.svd_up = svd_up self.svd_down = svd_down del self.sdnq_lora_svd_stash + lora_stack.drop(getattr(self, 'network_layer_name', None)) # a selection schedule must not outlive the segments it points into return True @@ -214,11 +215,23 @@ def apply_factors(self, network_layer_name, wanted_names): def append_factors(self, ups, downs): - """Concatenate ``[out, r]`` / ``[r, in]`` factor pairs onto the layer's svd channel and stash the originals.""" + """Concatenate ``[out, r]`` / ``[r, in]`` factor pairs onto the layer's svd channel and stash the originals. + + Returns the appended parts' rank ranges plus the transposed-layout flag; the + checkpoint's own factors occupy the range before the first entry and bucket + padding lands after the last, so the ranges stay valid on the live buffers. + """ deq = self.sdnq_dequantizer device = self.scale.device dtype = deq.result_dtype orig_up, orig_down = self.svd_up, self.svd_down + orig_rank = 0 + if orig_up is not None: + orig_rank = orig_up.shape[0] if deq.use_quantized_matmul else orig_up.shape[1] + segments, offset = [], orig_rank + for u in ups: + segments.append((offset, offset + u.shape[1])) + offset += u.shape[1] if deq.use_quantized_matmul: # matmul layout stores factors transposed: svd_up [r, out], svd_down [in, r] parts_up = ([orig_up.to(device=devices.device, dtype=dtype)] if orig_up is not None else []) + [u.t() for u in ups] @@ -240,6 +253,7 @@ def append_factors(self, ups, downs): self.sdnq_lora_svd_stash = (orig_up, orig_down) self.svd_up = torch.nn.Parameter(new_up.to(device=device), requires_grad=False) self.svd_down = torch.nn.Parameter(new_down.to(device=device), requires_grad=False) + return segments, deq.use_quantized_matmul def host_candidate(self, network_layer_name, wanted_names): @@ -434,6 +448,57 @@ def truncate_delta(self, D, dtype): return up_h, down_h, energy, rms is not None +def apply_select(self, network_layer_name, per_net, wanted_names): + """Attach two networks' contributions as separate side-channel segments for per-layer selection. + + Factorable members ride exactly; the rest host as their own truncated svd + with per-net cache entries. Segment ranges and selection scores register + with ``lora_stack``; the flip schedule executes from the step callback. + Returns None when the pair cannot ride the channel; the caller falls back. + """ + from sdnq.quant_utils import rotate_hadamard + deq = self.sdnq_dequantizer + changed = remove_factors(self) + if wanted_names == (): + return changed + if per_net is None or len(per_net) != 2: + return None + dtype = deq.result_dtype + lora_factor_cache.begin_pass(wanted_names) + pairs, ranks = [], [] + for i, (net_name, D) in enumerate(per_net): + if D is None or D.ndim != 2 or tuple(D.shape) != tuple(deq.original_shape): + return None + net = next((n for n in l.loaded_networks if n.name == net_name), None) + module = net.modules.get(network_layer_name, None) if net is not None else None + if module is None: + return None + ranks.append(int(getattr(module, 'dim', 0) or 0) or min(int(shared.opts.lora_sdnq_host_rank), *deq.original_shape)) + factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) + if factors is not None: + up_i, down_i = factors + if deq.use_hadamard: + down_i = rotate_hadamard(down_i.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype) + else: + key = f'{network_layer_name}#{i}' + cached = lora_factor_cache.fetch(key) + if cached is not None: + up_i, down_i = cached[0].to(device=devices.device, dtype=dtype), cached[1].to(device=devices.device, dtype=dtype) + hosted_layers.append((key, cached[2], cached[3])) + else: + up_i, down_i, energy, calibrated = truncate_delta(self, D.detach().to(devices.device, torch.float32), dtype) + up_i, down_i = lora_factor_cache.store(key, up_i, down_i, energy, calibrated, float(D.detach().float().square().mean().sqrt())) + hosted_layers.append((key, energy, calibrated)) + pairs.append((up_i, down_i)) + d0 = per_net[0][1].detach().to(devices.device, torch.float32) + d1 = per_net[1][1].detach().to(devices.device, torch.float32) + scores, abs_sums = lora_stack.score_pair(d0, d1, ranks[0], ranks[1]) + del d0, d1 + segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]]) + lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums) + return True + + def note_fallback(self, network_layer_name): """Record a quantized layer taking the requantize path (summary-logged per pass); layers the routing rule sent there are counted apart.""" if getattr(self, 'sdnq_dequantizer', None) is not None and network_layer_name not in routed_layers: diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index 224505cab..313038b39 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -30,7 +30,7 @@ KLORA_BETA = 0.5 # the paper's fixed ramp offset; only the slope is user-tunable ROW_CHUNK = 512 # fp32 interiors run in first-dim slices; also fixes the DARE draw sequence SAMPLE_CAP = 1 << 22 # strided subsample bound for magnitude quantiles (full-size quantile exceeds torch limits) -state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'total_steps': 0, 'finalized': False} +state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_num': 0.0, 'gamma_den': 0.0, 'total_steps': 0, 'finalized': False} warned: set = set() @@ -134,6 +134,45 @@ def combine(named_deltas, layer_name): return result.to(out_dtype) +def score_pair(d0, d1, rank0, rank1): + """Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance.""" + abs_sums = (float(d0.abs().sum()), float(d1.abs().sum())) + if mode() == 'klora': + k = max(1, int(rank0) * int(rank1)) + s0 = float(torch.topk(d0.abs().flatten(), min(k, d0.numel()), sorted=False).values.sum()) + s1 = float(torch.topk(d1.abs().flatten(), min(k, d1.numel()), sorted=False).values.sum()) + else: + s0 = float(d0.float().square().sum()) + s1 = float(d1.float().square().sum()) + return (s0, s1), abs_sums + + +def register_weight_pair(layer_name, module, per_net): + """Score and register a weight-kind selection pair; True when the layer is scheduled.""" + from modules.lora import lora_common as l + if per_net is None or len(per_net) != 2: + return False + ranks, names = [], [] + for net_name, d in per_net: + if d is None: + return False + net = next((n for n in l.loaded_networks if n.name == net_name), None) + net_module = net.modules.get(layer_name, None) if net is not None else None + if net_module is None: + return False + names.append(net_name) + ranks.append(int(getattr(net_module, 'dim', 0) or 0) or 64) + scores, abs_sums = score_pair(per_net[0][1].float(), per_net[1][1].float(), ranks[0], ranks[1]) + register(layer_name, module, 'weight', scores, nets=tuple(names), abs_sums=abs_sums) + return True + + +def drop(layer_name): + """Forget a layer's selection entry (its factors were removed or restored).""" + if layer_name is not None and state['entries'].pop(layer_name, None) is not None: + state['finalized'] = False + + def score_topk(up, down, k): """K-LoRA layer score: sum of the top-K absolute delta entries (one dense materialization).""" d = (up.to(torch.float32) @ down.to(torch.float32)).abs().flatten() @@ -152,22 +191,28 @@ def clear(): state['entries'] = {} state['flips'] = {} state['gamma'] = 1.0 + state['gamma_num'] = 0.0 + state['gamma_den'] = 0.0 state['total_steps'] = 0 state['finalized'] = False -def register(layer_name, module, kind, segments, scores, factors=None): - """Record a select-mode layer: its two segments (or bf16 factor pairs) and static scores. +def register(layer_name, module, kind, scores, segments=None, nets=None, abs_sums=None): + """Record a select-mode layer for schedule finalization. - kind 'factor': segments = [(start, stop), (start, stop)] column ranges in svd_up/svd_down - with the transposed-layout flag appended; stashes both segments' values for flips. - kind 'weight': factors = [(up0, down0), (up1, down1)] kept for recompute-from-backup. + kind 'factor': segments = ((s0, s1), (t0, t1), transposed) column ranges on the svd + channel; both segments' pristine values are stashed for flips. kind 'weight': nets = + the two network names; the winner delta is recomputed from the layer backup at + selection time. abs_sums feeds the global magnitude balance (klora gamma). """ - entry = {'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'factors': factors, 'stash': None} + entry = {'layer': layer_name, 'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'nets': nets, 'stash': None} if kind == 'factor': (s0, s1), (t0, t1), transposed = segments up = module.svd_up.data entry['stash'] = (segment_view(up, s0, s1, transposed).clone(), segment_view(up, t0, t1, transposed).clone()) + if abs_sums is not None: + state['gamma_num'] += abs_sums[0] + state['gamma_den'] += abs_sums[1] state['entries'][layer_name] = entry state['finalized'] = False @@ -196,6 +241,7 @@ def layer_flip_step(scores, total_steps): def finalize(total_steps): """Build the inverted flip map for the pass; select-mode layers start at their step-0 winner.""" state['total_steps'] = int(total_steps) + state['gamma'] = (state['gamma_num'] / state['gamma_den']) if state['gamma_den'] > 0 else 1.0 state['flips'] = {} for layer_name, entry in state['entries'].items(): flip_at = layer_flip_step(entry['scores'], state['total_steps']) @@ -208,7 +254,7 @@ def finalize(total_steps): def reset(total_steps): """Per-pass reset from set_callbacks_p: restore initial selections and reschedule for this pass's step count.""" - if mode() not in SELECT_MODES or not state['entries']: + if mode() not in SELECT_MODES or not state['entries'] or int(total_steps) <= 0: return finalize(total_steps) @@ -231,20 +277,25 @@ def apply_selection(layer_name, entry, winner): if entry['kind'] == 'factor': (s0, s1), (t0, t1), transposed = entry['segments'] up = module.svd_up.data - keep, drop = ((t0, t1), (s0, s1)) if winner == 1 else ((s0, s1), (t0, t1)) + keep_seg, drop_seg = ((t0, t1), (s0, s1)) if winner == 1 else ((s0, s1), (t0, t1)) stash = entry['stash'][winner] - segment_view(up, keep[0], keep[1], transposed).copy_(stash.to(device=up.device, dtype=up.dtype)) - segment_view(up, drop[0], drop[1], transposed).zero_() + segment_view(up, keep_seg[0], keep_seg[1], transposed).copy_(stash.to(device=up.device, dtype=up.dtype)) + segment_view(up, drop_seg[0], drop_seg[1], transposed).zero_() else: weight_selection(module, entry, winner) def weight_selection(module, entry, winner): + from modules.lora import lora_common as l + from modules.lora.lora_apply import network_apply_weights backup = getattr(module, 'network_weights_backup', None) if not isinstance(backup, torch.Tensor): # fuse mode keeps a bool sentinel, not a pristine copy warn_once('select-nobackup', 'Network stack: select flip skipped, no weight backup') return - up, down = entry['factors'][winner] - weight = backup.to(device=module.weight.device, dtype=torch.float32) - delta = up.to(device=module.weight.device, dtype=torch.float32) @ down.to(device=module.weight.device, dtype=torch.float32) - module.weight.data.copy_((weight + delta.reshape(weight.shape)).to(module.weight.dtype)) + net = next((n for n in l.loaded_networks if n.name == entry['nets'][winner]), None) + net_module = net.modules.get(entry['layer'], None) if net is not None else None + if net_module is None: + return + device = module.weight.device + updown = net_module.calc_updown(backup.to(device))[0] + network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 58548b82a..59a66a6e0 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -91,6 +91,7 @@ def network_activate(include=None, exclude=None): with devices.inference_context(), pbar: wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () stack_sig = lora_stack.signature() + lora_sdnq.signature() # tracked beside network_current_names so settings-only stack or mechanism changes re-apply + select_active = lora_stack.active_select(len(l.loaded_networks)) applied_layers.clear() lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind lora_sdnq.hosted_layers.clear() @@ -108,6 +109,37 @@ def network_activate(include=None, exclude=None): if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.loaded_networks): device = group_offload_strip(sd_model, component, group_stripped) calced = False # tracks whether this iteration assembled the delta, so the fallthrough reuses it instead of recomputing + if select_active and component_wanted and not network_layer_name.startswith('lora_te'): + if lora_sdnq.host_candidate(module, network_layer_name, component_wanted): # sub-8-bit SDNQ pairs ride the channel as separate segments + weights_backup = getattr(module, "network_weights_backup", None) + if weights_backup is not None and not isinstance(weights_backup, bool): + network_apply_weights(module, None, None, device=device) + per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) + if sel_bias is None: + applied = lora_sdnq.apply_select(module, network_layer_name, per_net, component_wanted) + if applied is not None: + if applied and component_wanted: + applied_layers.append(network_layer_name) + applied_weight += 1 + module.network_current_names = component_wanted + module.network_current_stack = stack_sig + if task is not None: + pbar.update(task, advance=1) + continue + else: # other layers select by recomputing the winner from the pristine backup at schedule time + backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse) + weights_backup = getattr(module, "network_weights_backup", None) + if weights_backup is not None and not isinstance(weights_backup, bool): + per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) + if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net): + network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner + applied_layers.append(network_layer_name) + applied_weight += 1 + module.network_current_names = component_wanted + module.network_current_stack = stack_sig + if task is not None: + pbar.update(task, advance=1) + continue if lora_sdnq.factor_candidate(module, network_layer_name, component_wanted): weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is not None and not isinstance(weights_backup, bool): @@ -153,6 +185,7 @@ def network_activate(include=None, exclude=None): continue backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse) if not component_wanted: + lora_stack.drop(network_layer_name) # a restored layer must leave the selection schedule weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is None or isinstance(weights_backup, bool): # fuse mode has no tensor backup, restore stays with network_deactivate if task is not None: diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index 9b8728695..1cc11dde5 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -17,6 +17,8 @@ def set_callbacks_p(processing): global p, warned # pylint: disable=global-statement p = processing warned = False + from modules.lora import lora_stack + lora_stack.reset(int(getattr(processing, 'steps', 0) or 0)) # per-pass: restore initial selections and reschedule flips before any step runs def prompt_callback(step, kwargs): @@ -37,6 +39,8 @@ def prompt_callback(step, kwargs): def diffusers_callback_legacy(step: int, timestep: int, latents: torch.FloatTensor | np.ndarray): if p is None: return + from modules.lora import lora_stack + lora_stack.on_step(step) if isinstance(latents, np.ndarray): # latents from Onnx pipelines is ndarray. latents = torch.from_numpy(latents) shared.state.sampling_step = step @@ -56,6 +60,8 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No if kwargs is None: kwargs = {} t0 = time.time() + from modules.lora import lora_stack + lora_stack.on_step(step) if shared.opts.torch_sync: if devices.backend == "ipex": diff --git a/scripts/xyz/xyz_grid_classes.py b/scripts/xyz/xyz_grid_classes.py index 2da29d7b3..198d93c91 100644 --- a/scripts/xyz/xyz_grid_classes.py +++ b/scripts/xyz/xyz_grid_classes.py @@ -104,6 +104,10 @@ class SharedSettingsStackHelper(): todo_ratio = None teacache_thresh = None extra_networks_default_multiplier = None + lora_stack_mode = None + lora_stack_density = None + lora_stack_alpha = None + lora_stack_discrepancy = None disable_apply_metadata = None disable_apply_params = None sdnq_quant_mode = None @@ -136,6 +140,10 @@ class SharedSettingsStackHelper(): self.sd_unet = shared.opts.sd_unet self.sd_text_encoder = shared.opts.sd_text_encoder self.extra_networks_default_multiplier = shared.opts.extra_networks_default_multiplier + self.lora_stack_mode = shared.opts.lora_stack_mode + self.lora_stack_density = shared.opts.lora_stack_density + self.lora_stack_alpha = shared.opts.lora_stack_alpha + self.lora_stack_discrepancy = shared.opts.lora_stack_discrepancy self.teacache_thresh = shared.opts.teacache_thresh self.disable_apply_metadata = shared.opts.disable_apply_metadata self.disable_apply_params = shared.opts.disable_apply_params @@ -148,6 +156,10 @@ class SharedSettingsStackHelper(): shared.opts.data["disable_apply_metadata"] = self.disable_apply_metadata shared.opts.data["disable_apply_params"] = self.disable_apply_params shared.opts.data["extra_networks_default_multiplier"] = self.extra_networks_default_multiplier + shared.opts.data["lora_stack_mode"] = self.lora_stack_mode + shared.opts.data["lora_stack_density"] = self.lora_stack_density + shared.opts.data["lora_stack_alpha"] = self.lora_stack_alpha + shared.opts.data["lora_stack_discrepancy"] = self.lora_stack_discrepancy shared.opts.data["prompt_attention"] = self.prompt_attention shared.opts.data["schedulers_solver_order"] = self.schedulers_solver_order shared.opts.data["schedulers_sigma_adjust"] = self.schedulers_sigma_adjust @@ -205,6 +217,10 @@ axis_options = [ AxisOption("[Prompt] Prompt parser", str, apply_setting("prompt_attention"), choices=lambda: ["native", "compel", "xhinker", "a1111", "fixed"]), AxisOption("[Network] LoRA", str, apply_lora, cost=0.5, choices=list_lora), AxisOption("[Network] LoRA strength", float, apply_lora_strength, cost=0.6), + AxisOption("[Network] LoRA stack mode", str, apply_setting("lora_stack_mode"), cost=0.6, choices=lambda: ["sum", "ties", "dare_ties", "dare_linear", "magnitude_prune", "klora", "estlora"]), + AxisOption("[Network] LoRA stack density", float, apply_setting("lora_stack_density"), cost=0.6), + AxisOption("[Network] LoRA stack ramp", float, apply_setting("lora_stack_alpha"), cost=0.6), + AxisOption("[Network] LoRA stack discrepancy", float, apply_setting("lora_stack_discrepancy"), cost=0.6), AxisOption("[Network] Styles", str, apply_styles, choices=lambda: [s.name for s in shared.prompt_styles.styles.values()]), AxisOption("[Param] Width", int, apply_field("width")), AxisOption("[Param] Height", int, apply_field("height")), diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 1949fd6a3..b1e1881d4 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -1562,6 +1562,228 @@ def test_sum_mode_keeps_exact_stacking(): return True +CAT_SELECT = category('stack-select') + + +@contextmanager +def select_mode(name, alpha=None, disc=None): + old = {k: getattr(shared.opts, k, None) for k in ('lora_stack_mode', 'lora_stack_alpha', 'lora_stack_discrepancy')} + shared.opts.lora_stack_mode = name + if alpha is not None: + shared.opts.lora_stack_alpha = alpha + if disc is not None: + shared.opts.lora_stack_discrepancy = disc + lora_stack.clear() + lora_stack.warned.clear() + try: + yield + finally: + for k, v in old.items(): + setattr(shared.opts, k, v) + lora_stack.clear() + + +def select_pair(layer, seed0=41, seed1=42, scale1=1.0): + A1, B1, D1 = make_delta(seed=seed0, sigma=1e-2) + A2, B2, D2 = make_delta(seed=seed1, sigma=1e-2) + if scale1 != 1.0: + A2, D2 = A2 * scale1, D2 * scale1 + n1 = make_net('subject', layer, A1, B1) + n2 = make_net('style', layer, A2, B2) + return n1, n2, D1, D2 + + +def test_select_flip_schedule_end_to_end(): + layer = build_layer('uint4') + n1, n2, D1, D2 = select_pair(layer) + with mock_model(lin=layer), select_mode('klora', alpha=1.5): + Wdq0 = dq(layer) + activate(n1, n2) + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None and entry['kind'] == 'factor', 'a factorable pair must register factor segments' + assert entry['segments'][0] == (0, 8) and entry['segments'][1] == (8, 16), f'segments {entry["segments"]}' + total = 20 + lora_stack.reset(total) + flips = [s for s, layers in lora_stack.state['flips'].items() for _ in layers] + assert len(flips) <= 1, 'a monotone ramp allows at most one flip per layer' + eff0 = dq(layer) - Wdq0 + winner0 = 0 if rho_of(eff0, D1) > rho_of(eff0, D2) else 1 + for s in range(total): + lora_stack.on_step(s) + eff1 = dq(layer) - Wdq0 + if flips: + assert rho_of(eff1, D2) > 0.99, 'after the flip the style delta must be selected' + assert rho_of(eff0, D1) > 0.99, 'before the flip the subject delta must be selected' + else: + assert rho_of(eff1, [D1, D2][winner0]) > 0.99 + activate() + assert torch.equal(dq(layer), Wdq0), 'removal from an end-of-schedule state must restore bit-exact' + return True + + +def test_select_initial_style_when_ramp_starts_won(): + layer = build_layer('uint4') + n1, n2, _D1, D2 = select_pair(seed0=43, seed1=44, scale1=8.0, layer=layer) # style delta dominates + with mock_model(lin=layer), select_mode('estlora', alpha=1.0, disc=0.5): + Wdq0 = dq(layer) + activate(n1, n2) + lora_stack.reset(20) + eff = dq(layer) - Wdq0 + assert rho_of(eff, D2) > 0.99, 'a layer whose style side wins at step 0 must start style-selected' + return True + + +def test_select_flip_is_inplace_and_shape_stable(): + layer = build_layer('uint4') + n1, n2, _D1, _D2 = select_pair(layer, seed0=45, seed1=46) + with mock_model(lin=layer), select_mode('klora'): + activate(n1, n2) + param_id = id(layer.svd_up) + shape = tuple(layer.svd_up.shape) + lora_stack.reset(20) + entry = lora_stack.state['entries']['lora_transformer_test'] + (s0, s1), (t0, t1), transposed = entry['segments'] + zeroed = lora_stack.segment_view(layer.svd_up.data, t0, t1, transposed) + kept = lora_stack.segment_view(layer.svd_up.data, s0, s1, transposed) + assert float(zeroed.abs().sum()) == 0.0 or float(kept.abs().sum()) == 0.0, 'exactly one segment must be zeroed initially' + for s in range(20): + lora_stack.on_step(s) + assert id(layer.svd_up) == param_id and tuple(layer.svd_up.shape) == shape, 'flips must mutate in place, never reassign' + return True + + +def test_select_matmul_transposed_layout(): + layer = build_layer('uint4', use_quantized_matmul=True) + n1, n2, D1, D2 = select_pair(layer, seed0=47, seed1=48) + with mock_model(lin=layer), select_mode('klora'): + Wdq0 = dq(layer) + activate(n1, n2) + entry = lora_stack.state['entries']['lora_transformer_test'] + assert entry['segments'][2] is True, 'quantized-matmul layout must register as transposed' + lora_stack.reset(20) + eff = dq(layer) - Wdq0 + assert max(rho_of(eff, D1), rho_of(eff, D2)) > 0.99, 'initial selection must realize one delta exactly' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_select_per_net_hosted_pair(): + layer = build_layer('uint4') + torch.manual_seed(49) + Dd1 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 # rank inside the host cap so truncation is near-lossless + Dd2 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + n1 = make_dense_net('lk1', layer, Dd1) + n2 = make_dense_net('lk2', layer, Dd2) + with host_rank(32), mock_model(lin=layer), select_mode('klora'): + Wdq0 = dq(layer) + activate(n1, n2) + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None, 'non-factorable pairs must register through per-net hosting' + assert entry['segments'][0] == (0, 24) and entry['segments'][1] == (24, 48), f'segments {entry["segments"]}' # hosting stores the effective rank (24), not the cap + lora_stack.reset(20) + eff = dq(layer) - Wdq0 + best = max(rho_of(eff, Dd1), rho_of(eff, Dd2)) + assert best > 0.9, f'initial selection must realize one hosted delta, rho={best:.3f}' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_select_reset_restores_initial_state(): + layer = build_layer('uint4') + n1, n2, _D1, _D2 = select_pair(layer, seed0=51, seed1=52) + with mock_model(lin=layer), select_mode('klora'): + activate(n1, n2) + lora_stack.reset(20) + initial = dq(layer) + for s in range(20): + lora_stack.on_step(s) + lora_stack.reset(20) + assert torch.equal(dq(layer), initial), 'a fresh pass must restore the initial selection without re-activation' + return True + + +def test_select_deactivate_from_midflip(): + layer = build_layer('uint4') + n1, n2, _D1, _D2 = select_pair(layer, seed0=53, seed1=54) + with mock_model(lin=layer), select_mode('klora'): + Wdq0 = dq(layer) + activate(n1, n2) + lora_stack.reset(20) + for s in range(10): + lora_stack.on_step(s) + activate() + assert torch.equal(dq(layer), Wdq0), 'removal mid-schedule must restore bit-exact' + assert not lora_stack.state['entries'], 'removal must drop the selection entry' + return True + + +def test_select_requires_exactly_two_nets(): + layer = build_layer('uint4') + A3, B3, _D3 = make_delta(seed=55) + n1, n2, _D1, _D2 = select_pair(layer, seed0=56, seed1=57) + n3 = make_net('third', layer, A3, B3) + with mock_model(lin=layer), select_mode('klora'): + activate(n1, n2, n3) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'three nets must fall back to the exact concat path' + assert not lora_stack.state['entries'], 'no selection entries outside the two-net case' + activate() + return True + + +def test_select_gated_off_when_compiled(): + layer = build_layer('uint4') + n1, n2, _D1, _D2 = select_pair(layer, seed0=58, seed1=59) + old_compile = getattr(shared.opts, 'cuda_compile', None) + try: + shared.opts.cuda_compile = ['Model'] + with mock_model(lin=layer), select_mode('klora'): + activate(n1, n2) + assert not lora_stack.state['entries'], 'select must gate off under model compile' + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'gated select behaves as sum' + activate() + finally: + shared.opts.cuda_compile = old_compile + return True + + +def test_est_energy_matches_full_frobenius(): + torch.manual_seed(60) + up = torch.randn(64, 8, device=DEVICE) + down = torch.randn(8, 96, device=DEVICE) + gram = lora_stack.score_energy(up, down) + full = float((up @ down).square().sum()) + assert abs(gram - full) / full < 1e-5, f'{gram} vs {full}' + return True + + +def test_select_weight_kind_plain_layer(): + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device=DEVICE) + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.02) + lin.network_layer_name = 'lora_transformer_plain' + lin.network_current_names = () + A1, B1, D1 = make_delta(seed=61, sigma=1e-2) + A2, B2, D2 = make_delta(seed=62, sigma=1e-2) + n1 = make_net('w1', lin, A1, B1) + n2 = make_net('w2', lin, A2, B2) + W0 = lin.weight.detach().float().clone() + with mock_model(lin=lin), select_mode('klora'): + activate(n1, n2) + entry = lora_stack.state['entries'].get('lora_transformer_plain') + assert entry is not None and entry['kind'] == 'weight', 'plain layers must register weight-kind selection' + assert torch.equal(lin.weight.detach().float(), W0), 'weights stay pristine until the schedule applies a winner' + lora_stack.reset(20) + eff = lin.weight.detach().float() - W0 + assert max(rho_of(eff, D1), rho_of(eff, D2)) > 0.95, 'initial selection must apply one delta from backup' + for s in range(20): + lora_stack.on_step(s) + activate() + assert torch.equal(lin.weight.detach().float(), W0), 'restore-only pass must return the pristine weight' + return True + + CAT_COMPILE = category('compile') @@ -1746,6 +1968,12 @@ def run_tests(): test_magnitude_prune_keeps_top_density, test_dense_two_plain_loras_hosted_not_summed, test_single_net_ignores_dense_mode, test_te_layer_stays_plain_sum, test_sum_mode_keeps_exact_stacking]: run_test(CAT_STACK, fn) + log.warning('=== Stack modes: select ===') + for fn in [test_select_flip_schedule_end_to_end, test_select_initial_style_when_ramp_starts_won, test_select_flip_is_inplace_and_shape_stable, + test_select_matmul_transposed_layout, test_select_per_net_hosted_pair, test_select_reset_restores_initial_state, + test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled, + test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer]: + run_test(CAT_SELECT, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: run_test(CAT_COMPILE, fn) From e2202bfbdfc8a7b6592287acac9d74792e1a3f0c Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 03:32:53 +0100 Subject: [PATCH 03/23] feat(lora): log exact side-channel applies The exact factor path was the only apply route with no log line; its success read as silence. Track layers taking it beside the hosted and fallback lists and report all three as key=value apply lines (apply=exact/hosted/requantize); the stack fallback notices use the same form. The suite pins its stack-mode baseline to sum so a mode left set in user config cannot reroute tests that assume plain summation. --- modules/lora/extra_networks_lora.py | 2 +- modules/lora/lora_sdnq.py | 10 ++++++++-- modules/lora/lora_stack.py | 6 +++--- modules/lora/networks.py | 1 + test/test-sdnq-lora-factors.py | 1 + 5 files changed, 14 insertions(+), 6 deletions(-) diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index c81b94072..bd44ba7f5 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -227,7 +227,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): load_method, load_reason = lora_overrides.get_method() from modules.lora import lora_stack if load_method != 'native' and lora_stack.mode() != 'sum': - lora_stack.warn_once(f'method-{load_method}', f'Network stack: mode={lora_stack.mode()} method={load_method} unsupported, using sum') + lora_stack.warn_once(f'method-{load_method}', f'Network stack: mode={lora_stack.mode()} method={load_method} fallback=sum') if debug: import sys fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index b1ef10d45..df2e7d2b0 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -55,6 +55,7 @@ from modules.logger import log fallback_layers: list[str] = [] hosted_layers: list[tuple[str, float, bool]] = [] hosted_ranks: list[int] = [] +factor_layers: list[str] = [] routed_layers: list[str] = [] REQUANT_RATIO = 0.30 # delta rms over mean grid step above which requantize can retain the delta @@ -211,6 +212,7 @@ def apply_factors(self, network_layer_name, wanted_names): if not ups: return changed append_factors(self, ups, downs) + factor_layers.append(network_layer_name) return True @@ -496,6 +498,7 @@ def apply_select(self, network_layer_name, per_net, wanted_names): del d0, d1 segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]]) lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums) + factor_layers.append(network_layer_name) return True @@ -509,6 +512,9 @@ def report_fallbacks(): hits, misses = lora_factor_cache.flush() if hits > 0 or misses > 0: log.info(f'Network load: type=LoRA quant=sdnq cache hits={hits} misses={misses}') + if len(factor_layers) > 0: + log.info(f'Network load: type=LoRA quant=sdnq apply=exact layers={len(factor_layers)}') + factor_layers.clear() if len(hosted_layers) > 0: energies = sorted(e for _name, e, _c in hosted_layers) median = energies[len(energies) // 2] @@ -517,7 +523,7 @@ def report_fallbacks(): if len(hosted_ranks) > 0 and min(hosted_ranks) < int(shared.opts.lora_sdnq_host_rank): rs = sorted(hosted_ranks) ranks = f' k={rs[0]}-{rs[len(rs) // 2]}-{rs[-1]}' # realized rank spread; shown only when a spectrum collapsed below the cap - log.info(f'Network load: type=LoRA quant=sdnq hosted={len(hosted_layers)} rank={int(shared.opts.lora_sdnq_host_rank)}{ranks}{f" calib={calibrated}" if calibrated else ""} energy={median:.2f} min={energies[0]:.2f} non-factorable networks hosted on the svd side-channel') + log.info(f'Network load: type=LoRA quant=sdnq apply=hosted layers={len(hosted_layers)} rank={int(shared.opts.lora_sdnq_host_rank)}{ranks}{f" calib={calibrated}" if calibrated else ""} energy={median:.2f} min={energies[0]:.2f}') if l.debug: log.debug(f'Network load: type=LoRA quant=sdnq hosted={[(n, round(e, 3)) for n, e, _c in hosted_layers[:8]]}{"..." if len(hosted_layers) > 8 else ""}') hosted_layers.clear() @@ -529,7 +535,7 @@ def report_fallbacks(): routed_layers.clear() if len(fallback_layers) > 0: if enabled(): - log.warning(f'Network load: type=LoRA quant=sdnq layers={len(fallback_layers)} non-factorable networks requantized in place (reduced fidelity on quantized weights)') + log.warning(f'Network load: type=LoRA quant=sdnq apply=requantize layers={len(fallback_layers)} fidelity=reduced') else: log.info(f'Network load: type=LoRA quant=sdnq apply=requantize layers={len(fallback_layers)} reason=setting') if l.debug: diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index 313038b39..b1d889bda 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -78,10 +78,10 @@ def active_select(n_loaded): if m not in SELECT_MODES: return False if n_loaded != 2: - warn_once('select-count', f'Network stack: mode={m} networks={n_loaded} requires exactly 2, using sum') + warn_once('select-count', f'Network stack: mode={m} networks={n_loaded} required=2 fallback=sum') return False if select_blocked(): - warn_once('select-compile', f'Network stack: mode={m} disabled with model compile, using sum') + warn_once('select-compile', f'Network stack: mode={m} compile=model fallback=sum') return False return True @@ -290,7 +290,7 @@ def weight_selection(module, entry, winner): from modules.lora.lora_apply import network_apply_weights backup = getattr(module, 'network_weights_backup', None) if not isinstance(backup, torch.Tensor): # fuse mode keeps a bool sentinel, not a pristine copy - warn_once('select-nobackup', 'Network stack: select flip skipped, no weight backup') + warn_once('select-nobackup', 'Network stack: flip=skipped backup=none') return net = next((n for n in l.loaded_networks if n.name == entry['nets'][winner]), None) net_module = net.modules.get(entry['layer'], None) if net is not None else None diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 59a66a6e0..be543d2b2 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -95,6 +95,7 @@ def network_activate(include=None, exclude=None): applied_layers.clear() lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind lora_sdnq.hosted_layers.clear() + lora_sdnq.factor_layers.clear() backup_size = 0 for component in modules.keys(): component_wanted = wanted_names if component in components else () diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index b1e1881d4..7259cca26 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -82,6 +82,7 @@ from sdnq.quantizer import sdnq_quantize_layer, SDNQConfig # pylint: disable=wr DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') OUT_F, IN_F, RANK = 512, 512, 8 +shared.opts.lora_stack_mode = 'sum' # suite baseline regardless of user config; stack tests set modes via their own context managers results: dict[str, dict] = {} From 734215cc6dca8895b078124b06fd0bffa8979018 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 03:41:28 +0100 Subject: [PATCH 04/23] fix(lora): keep select stack modes dormant without a qualifying pair A select mode forced backup mode whenever it was merely set, so a leftover setting changed behavior for ordinary single-network loads. The fuse gate now engages only when the loaded set could actually select (exactly two networks, compile permitting) or while selection segments are still live on model layers. Re-application drops any stale per-layer schedule so a later pass reset can never replay an old winner over freshly applied weights. Fallback notices log per activation instead of once per session; only the in-loop flip notice stays latched. --- modules/lora/extra_networks_lora.py | 2 +- modules/lora/lora_overrides.py | 5 ++-- modules/lora/lora_stack.py | 14 +++++++++-- modules/lora/networks.py | 1 + test/test-sdnq-lora-factors.py | 37 +++++++++++++++++++++++++++++ 5 files changed, 54 insertions(+), 5 deletions(-) diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index bd44ba7f5..b08dbfaf8 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -227,7 +227,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): load_method, load_reason = lora_overrides.get_method() from modules.lora import lora_stack if load_method != 'native' and lora_stack.mode() != 'sum': - lora_stack.warn_once(f'method-{load_method}', f'Network stack: mode={lora_stack.mode()} method={load_method} fallback=sum') + log.warning(f'Network stack: mode={lora_stack.mode()} method={load_method} fallback=sum') if debug: import sys fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access diff --git a/modules/lora/lora_overrides.py b/modules/lora/lora_overrides.py index 078b285b2..305d1e0e6 100644 --- a/modules/lora/lora_overrides.py +++ b/modules/lora/lora_overrides.py @@ -112,9 +112,10 @@ def disable_fuse(): round-trips it through its storage format. On quantized weights that is a dequantize-add-requantize cycle per network swap whose error compounds. """ + from modules.lora import lora_common as l from modules.lora import lora_stack - if lora_stack.mode() in lora_stack.SELECT_MODES: - return True # select stack modes flip per-layer winners against the pristine backup + if lora_stack.select_possible(len(l.loaded_networks)) or lora_stack.select_engaged(): + return True # select flips per-layer winners against the pristine backup; a dormant select mode leaves fuse alone sd_model = getattr(shared.sd_model, 'pipe', shared.sd_model) if is_quantized(sd_model): return True diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index b1d889bda..a9899b1d4 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -73,15 +73,25 @@ def active_dense(n_contrib): return mode() in DENSE_MODES and n_contrib >= 2 +def select_possible(n_loaded): + """True when the loaded set could engage a select mode; silent, for the fuse gate.""" + return mode() in SELECT_MODES and n_loaded == 2 and not select_blocked() + + +def select_engaged(): + """True while selection schedules are live on model layers.""" + return bool(state['entries']) + + def active_select(n_loaded): m = mode() if m not in SELECT_MODES: return False if n_loaded != 2: - warn_once('select-count', f'Network stack: mode={m} networks={n_loaded} required=2 fallback=sum') + log.warning(f'Network stack: mode={m} networks={n_loaded} required=2 fallback=sum') return False if select_blocked(): - warn_once('select-compile', f'Network stack: mode={m} compile=model fallback=sum') + log.warning(f'Network stack: mode={m} compile=model fallback=sum') return False return True diff --git a/modules/lora/networks.py b/modules/lora/networks.py index be543d2b2..6e442f5f3 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -107,6 +107,7 @@ def network_activate(include=None, exclude=None): if task is not None: pbar.update(task, advance=1) continue + lora_stack.drop(network_layer_name) # re-application invalidates any live selection schedule; the select branch re-registers if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.loaded_networks): device = group_offload_strip(sd_model, component, group_stripped) calced = False # tracks whether this iteration assembled the delta, so the fallthrough reuses it instead of recomputing diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 7259cca26..d48e660c1 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -1749,6 +1749,42 @@ def test_select_gated_off_when_compiled(): return True +def test_select_gate_dormant_without_pair(): + with select_mode('klora'): + assert lora_stack.select_possible(1) is False, 'a single network must leave the fuse gate alone' + assert lora_stack.select_possible(2) is True + assert lora_stack.select_possible(3) is False + old_compile = getattr(shared.opts, 'cuda_compile', None) + try: + shared.opts.cuda_compile = ['Model'] + assert lora_stack.select_possible(2) is False, 'compile block must keep the gate down' + finally: + shared.opts.cuda_compile = old_compile + assert lora_stack.select_engaged() is False + with select_mode('sum'): + assert lora_stack.select_possible(2) is False + return True + + +def test_stale_schedule_dropped_on_reapply(): + layer = build_layer('uint4') + n1, n2, D1, _D2 = select_pair(layer, seed0=64, seed1=65) + with mock_model(lin=layer), select_mode('klora'): + Wdq0 = dq(layer) + activate(n1, n2) + assert lora_stack.select_engaged(), 'the pair must register schedules' + lora_stack.reset(20) + activate(n1) # same mode still set, but a single net cannot select + assert not lora_stack.select_engaged(), 're-application must drop the stale schedule' + w_single = dq(layer) + assert rho_of(w_single - Wdq0, D1) > 0.99, 'the single net must apply exactly' + lora_stack.reset(20) # a later pass reset must find nothing to replay + assert torch.equal(dq(layer), w_single), 'a stale schedule must never overwrite a fresh apply' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + def test_est_energy_matches_full_frobenius(): torch.manual_seed(60) up = torch.randn(64, 8, device=DEVICE) @@ -1973,6 +2009,7 @@ def run_tests(): for fn in [test_select_flip_schedule_end_to_end, test_select_initial_style_when_ramp_starts_won, test_select_flip_is_inplace_and_shape_stable, test_select_matmul_transposed_layout, test_select_per_net_hosted_pair, test_select_reset_restores_initial_state, test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled, + test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===') From c42adcf6806bdba1d36cfb3a17511bea1d9709e8 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 04:01:36 +0100 Subject: [PATCH 05/23] docs(i18n): rework lora hosting and stack mode hints Restructure the seven hints to the settings pattern: lead definition, tradeoff, scope, special values and stated default, with value and cross-setting markup. Mode bullets and special values match loader behavior: sub-8-bit hosting gate, rank 0 fallback, ramp 0 freezing the balance. --- ui/locale/locale_en.json | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 522ae19e6..678dede3e 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -865,13 +865,13 @@ {"id":"","label":"LoRA native fuse with model","localized":"","hint":"Merge LoRA into the model for lower memory usage.

Warning: After removing or switching a LoRA, you may still see its style in generated images. To get a clean model, reload it from the model selector.","ui":"settings_lora"}, {"id":"","label":"LoRA diffusers fuse with model","localized":"","hint":"Merge LoRA into the model for lower memory usage and torch.compile compatibility.

Warning: After removing or switching a LoRA, you may still see its style in generated images. To get a clean model, reload it from the model selector.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized apply method","localized":"","hint":"How networks are applied to SDNQ-quantized model weights:
- exact: adapters are carried alongside the quantized weights at full precision; apply and removal are exact and the quantized weights are never modified. The carried factors take additional VRAM, growing with adapter rank, size and count
- requantize: adapters are merged into the quantized weights, matching the behavior of earlier releases. Uses no additional VRAM (a weight backup for network removal is held in system RAM); on models quantized below 8 bits rounding typically loses much of the adapter effect, with strong adapters retaining more

With requantize selected, the host rank, calibration and cache options below have no effect.

Default is exact.","ui":"settings_lora"}, - {"id":"","label":"LoRA quantized host rank","localized":"","hint":"Maximum rank used to carry non-factorable adapter types (LoKR, LoHA, OFT, DoRA) on the side-channel of SDNQ models quantized below 8 bits, where merging would erase most of the adapter. Higher values keep more of the adapter at proportionally higher memory cost. Set to 0 to disable and merge into the quantized weights instead.","ui":"settings_lora"}, - {"id":"","label":"LoRA quantized host calibration","localized":"","hint":"Collection of per-channel activation statistics during normal generations on models quantized below 8 bits, cached per checkpoint. When available, the statistics weight side-channel hosting of non-factorable adapter types toward the channels carrying the most activation energy, improving delivered fidelity at the same host rank.","ui":"settings_lora"}, - {"id":"","label":"LoRA quantized host cache","localized":"","hint":"Disk budget in GB for caching the side-channel factors computed when hosting non-factorable adapter types on quantized models. A cached configuration skips the truncation math on reapply; least recently used entries are evicted past the budget. Set to 0 to disable.","ui":"settings_lora"}, - {"id":"","label":"LoRA stack mode","localized":"","hint":"Combination rule applied when multiple LoRA networks target the same layer. sum adds all contributions. ties merges only elements whose signs agree after keeping the strongest ones; dare variants randomly drop elements and rescale the survivors; magnitude_prune keeps only the strongest elements of each network. klora and estlora select one of exactly two networks per layer, with the first network in the prompt treated as subject and the second as style, and the balance shifting from subject toward style over the sampling steps. Applied by the native loader only; other load methods behave as sum.","ui":"settings_lora"}, - {"id":"","label":"LoRA stack density","localized":"","hint":"Fraction of elements kept by the ties, dare and magnitude_prune stack modes. Lower values keep only the strongest contributions of each network.","ui":"settings_lora"}, - {"id":"","label":"LoRA stack ramp","localized":"","hint":"Steepness of the subject-to-style shift over the sampling steps for the klora and estlora stack modes. Higher values strengthen the late-step shift toward the style network.","ui":"settings_lora"}, - {"id":"","label":"LoRA stack discrepancy","localized":"","hint":"Stand-in for the estlora mode's data-derived style separation estimate. Higher values shift the overall balance toward the subject network.","ui":"settings_lora"}, + {"id":"","label":"LoRA quantized host rank","localized":"","hint":"Maximum rank used to carry adapter types that are not natively low-rank (LoKR, LoHA, OFT, DoRA) alongside the quantized weights instead of merging them in.
Higher values retain more of the adapter at proportionally more memory. Plain LoRA files are carried exactly at their own rank.

Applies only to SDNQ models quantized below 8 bits, where merging erases most of the adapter; at 8 bits and above merging retains it and hosting is skipped.

0 disables hosting and merges every adapter into the quantized weights.

Default is 256.","ui":"settings_lora"}, + {"id":"","label":"LoRA quantized host calibration","localized":"","hint":"Collects per-channel activation statistics from the model's own generations and uses them to focus hosted-adapter truncation on the channels with the strongest activations.
Statistics accumulate in the background on models quantized below 8 bits, persist per checkpoint, and raise delivered adapter fidelity at the same LoRA quantized host rank, most at low ranks.

Capture is skipped while the model is compiled; previously cached statistics still apply.

Enabled by default.","ui":"settings_lora"}, + {"id":"","label":"LoRA quantized host cache","localized":"","hint":"Disk space in GB for caching computed hosting factors.
A cached set skips the truncation math on the next load; least recently used entries are evicted once the budget is exceeded.

0 disables the cache.

Default is 10.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack mode","localized":"","hint":"How multiple networks targeting the same layer are combined:
- sum: adds all contributions
- ties: keeps each network's strongest elements and merges only where signs agree
- dare_ties: randomly drops elements, rescales the survivors, then merges where signs agree
- dare_linear: randomly drops elements, rescales the survivors and sums
- magnitude_prune: keeps each network's strongest elements and sums
- klora / estlora: assign each layer to one of exactly two networks, the first in the prompt as subject and the second as style, shifting from subject toward style over the sampling steps

Kept fractions are set by LoRA stack density; the subject-to-style shift by LoRA stack ramp and LoRA stack discrepancy.

Applies to the native load path; other load methods and text encoder networks always combine as sum. Selection modes fall back to sum unless exactly two networks are loaded, or when model compile is active.

Default is sum.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack density","localized":"","hint":"Fraction of elements each network keeps under the ties, dare_ties, dare_linear and magnitude_prune stack modes.
Lower values keep only the strongest contributions and reduce interference between networks at the cost of per-network detail. The dare variants drop at random and rescale the survivors to preserve expected strength.

Default is 0.5.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack ramp","localized":"","hint":"Slope of the subject-to-style shift across the sampling steps in the klora and estlora stack modes.
Higher values shift layers to the style network earlier and more broadly; lower values keep the subject network dominant for longer.

0 keeps the balance fixed for the whole generation.

Default is 1.5.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack discrepancy","localized":"","hint":"Stand-in for the measured style separation the estlora stack mode would otherwise derive from data.
Higher values keep layers with the subject network longer; lower values let the style network take layers earlier.

Applies only when LoRA stack mode is estlora.

Default is 0.5.","ui":"settings_lora"}, {"id":"","label":"LoRA auto-apply tags","localized":"","hint":"Automatically add trigger words/tags from LoRA metadata to your prompt.
Set to the number of tags to auto-apply, e.g., 3 = add top 3 trigger tags.
Set to 0 to disable, -1 to add all available tags.","ui":"settings_lora"}, {"id":"","label":"LoRA memory cache","localized":"","hint":"How many LoRAs to keep in network for future use before requiring reloading from storage","ui":"settings_lora"}, {"id":"","label":"LoRA add hash info to metadata","localized":"","hint":"Include LoRA file hashes in generated image metadata.
Useful for reproducibility and tracking which exact LoRA versions were used.","ui":"settings_lora"}, From 55c3eb1325598f0d081503c3cd0acf42c08c58bd Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 04:16:29 +0100 Subject: [PATCH 06/23] refactor(lora): remove the lora_apply_te setting Text encoder networks now apply unconditionally in the native path. - remove the option, the per-request parameter and the control threading - collapse activate_filtered into plain activate at all call sites - drop the toggle from the prompt embed cache key - register the retired key so existing configs load without warnings --- modules/control/run.py | 4 ++-- modules/detailer/detailer.py | 2 +- modules/extra_networks.py | 9 --------- modules/face/faceid.py | 2 +- modules/options_handler.py | 2 +- modules/processing_class.py | 2 -- modules/processing_diffusers.py | 4 ++-- modules/prompt_parser_diffusers.py | 5 +---- modules/ui_definitions.py | 1 - ui/locale/locale_en.json | 1 - 10 files changed, 8 insertions(+), 24 deletions(-) diff --git a/modules/control/run.py b/modules/control/run.py index f5dae8c2b..a64fabe60 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -403,7 +403,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg sequential_seed: bool | None = None, # prompt/attention overrides prompt_attention: str | None = None, prompt_mean_norm: bool | None = None, diffusers_zeros_prompt_pad: bool | None = None, - te_pooled_embeds: bool | None = None, lora_apply_te: bool | None = None, te_complex_human_instruction: str | None = None, te_use_mask: bool | None = None, + te_pooled_embeds: bool | None = None, te_complex_human_instruction: str | None = None, te_use_mask: bool | None = None, # generation modifier overrides (hijack) freeu_enabled: bool | None = None, freeu_b1: float | None = None, freeu_b2: float | None = None, freeu_s1: float | None = None, freeu_s2: float | None = None, hypertile_unet_enabled: bool | None = None, hypertile_hires_only: bool | None = None, hypertile_unet_tile: int | None = None, hypertile_unet_min_tile: int | None = None, @@ -586,7 +586,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg # prompt/attention overrides prompt_attention=prompt_attention, prompt_mean_norm=prompt_mean_norm, diffusers_zeros_prompt_pad=diffusers_zeros_prompt_pad, te_pooled_embeds=te_pooled_embeds, - lora_apply_te=lora_apply_te, te_complex_human_instruction=te_complex_human_instruction, te_use_mask=te_use_mask, + te_complex_human_instruction=te_complex_human_instruction, te_use_mask=te_use_mask, # generation modifier overrides (hijack) freeu_enabled=freeu_enabled, freeu_b1=freeu_b1, freeu_b2=freeu_b2, freeu_s1=freeu_s1, freeu_s2=freeu_s2, hypertile_unet_enabled=hypertile_unet_enabled, hypertile_hires_only=hypertile_hires_only, diff --git a/modules/detailer/detailer.py b/modules/detailer/detailer.py index 752e9fbdc..73bdd32f0 100644 --- a/modules/detailer/detailer.py +++ b/modules/detailer/detailer.py @@ -321,7 +321,7 @@ class Detailer(): pc.disable_extra_networks = True # disable processing_diffusers from handling network activation since its handled here network_same = len(p.network_data.values()) == len(pc.network_data.values()) and all(x == y for x, y in zip(p.network_data.values(), pc.network_data.values())) if not network_same: - extra_networks.activate_filtered(pc, pc.network_data) + extra_networks.activate(pc, pc.network_data) log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label}" score={item.score:.2f} seg={detailer_opt(p, "detailer_segmentation")} network={network_same} prompt="{pc.prompt}"') pc.init_images = [image] pc.image_mask = [item.mask] diff --git a/modules/extra_networks.py b/modules/extra_networks.py index 0737f68f8..89acf5a4e 100644 --- a/modules/extra_networks.py +++ b/modules/extra_networks.py @@ -121,15 +121,6 @@ def activate(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, p.network_data = extra_network_data -def activate_filtered(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, list[ExtraNetworkParams]] | None = None, step=0): - """activate with text encoder components gated on lora_apply_te; must run before prompt encode so te networks affect embeds""" - apply_te = getattr(p, 'lora_apply_te', None) - if apply_te is None: - apply_te = shared.opts.lora_apply_te - exclude = [] if apply_te else ['text_encoder', 'text_encoder_2', 'text_encoder_3'] - activate(p, extra_network_data, step=step, exclude=exclude) - - def deactivate(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, list[ExtraNetworkParams]] | None = None, force: bool | None = None): """call deactivate for extra networks in extra_network_data in specified order, then call deactivate for all remaining registered networks""" if p.disable_extra_networks: diff --git a/modules/face/faceid.py b/modules/face/faceid.py index 4c63bcd66..b2263d840 100644 --- a/modules/face/faceid.py +++ b/modules/face/faceid.py @@ -216,7 +216,7 @@ def face_id( p.subseeds = p.all_subseeds[n * p.batch_size:(n+1) * p.batch_size] p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts, p.network_data) - extra_networks.activate_filtered(p, p.network_data) + extra_networks.activate(p, p.network_data) ip_model_dict.update({ "prompt": p.prompts[0], "negative_prompt": p.negative_prompts[0], diff --git a/modules/options_handler.py b/modules/options_handler.py index 2400a2069..8b7818994 100644 --- a/modules/options_handler.py +++ b/modules/options_handler.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: import builtins cmd_opts = cmd_args.parse_args() -compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order', 'xformers_options'] +compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order', 'xformers_options', 'lora_apply_te'] removed_values = { # a stored choice that no longer exists is kept by validate, so it has to be rewritten or it selects nothing 'cross_attention_optimization': (['Batch matrix-matrix', 'Dynamic Attention BMM'], 'Scaled-Dot-Product'), } diff --git a/modules/processing_class.py b/modules/processing_class.py index b5f567379..624e874b7 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -254,7 +254,6 @@ class StableDiffusionProcessing: prompt_mean_norm: bool | None = None, diffusers_zeros_prompt_pad: bool | None = None, te_pooled_embeds: bool | None = None, - lora_apply_te: bool | None = None, te_complex_human_instruction: str | None = None, te_use_mask: bool | None = None, # generation modifier overrides (hijack) @@ -543,7 +542,6 @@ class StableDiffusionProcessing: self.prompt_mean_norm = prompt_mean_norm self.diffusers_zeros_prompt_pad = diffusers_zeros_prompt_pad self.te_pooled_embeds = te_pooled_embeds - self.lora_apply_te = lora_apply_te self.te_complex_human_instruction = te_complex_human_instruction self.te_use_mask = te_use_mask # generation modifier overrides (hijack) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 0f975ea72..a7cad820a 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -149,7 +149,7 @@ def process_base(p: processing.StableDiffusionProcessing): if 'detailer' in p.ops: desc = 'Detail' p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts, p.network_data) - extra_networks.activate_filtered(p) # networks must patch weights before prompt encode so te loras affect embeds + extra_networks.activate(p) # networks must patch weights before prompt encode so te loras affect embeds base_args = set_pipeline_args( p=p, model=shared.sd_model, @@ -313,7 +313,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output): prompts, p.network_data = extra_networks.parse_prompts(prompts) reset_prompts = True if reset_prompts or ('base' in p.skip): - extra_networks.activate_filtered(p) + extra_networks.activate(p) hires_args = set_pipeline_args( p=p, diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 0b073baa8..d02817733 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -130,11 +130,8 @@ class PromptEmbedder: # unpack EN data in case of TE LoRA en_data = p.network_data en_data = [idx.items for item in en_data.values() for idx in item] - apply_te = getattr(p, 'lora_apply_te', None) - if apply_te is None: - apply_te = shared.opts.lora_apply_te effective_batch = 1 if self.allsame else self.batchsize - key = str([self.prompts, self.negative_prompts, effective_batch, self.clip_skip, self.steps, en_data, apply_te]) + key = str([self.prompts, self.negative_prompts, effective_batch, self.clip_skip, self.steps, en_data]) item = cache.get(key) if not item: if not any(flatten(emb) for emb in [self.prompt_embeds, diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 93ebbb3aa..1808f62a7 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -699,7 +699,6 @@ def create_settings(cmd_opts): "lora_apply_tags": OptionInfo(0, "LoRA auto-apply tags", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}), "lora_apply_sep": OptionInfo("

Apply method

", "", gr.HTML), - "lora_apply_te": OptionInfo(False, "LoRA native apply to text encoder"), "lora_fuse_native": OptionInfo(True, "LoRA native fuse with model"), "lora_fuse_diffusers": OptionInfo(False, "LoRA diffusers fuse with model"), diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 678dede3e..2dd87c144 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -861,7 +861,6 @@ {"id":"","label":"Load custom Diffusers pipeline","localized":"","hint":"","ui":"settings_huggingface"}, {"id":"","label":"LoRA force reload always","localized":"","hint":"Forces LoRA networks to reload from storage on every generation, even if already cached.
Useful for debugging or when LoRA files are being modified externally.
Disable for normal use to benefit from caching.","ui":"settings_lora"}, {"id":"","label":"LoRA load using Diffusers method","localized":"","hint":"Alternative method uses diffusers built-in LoRA capabilities instead of native SD.Next implementation (may reduce LoRA compatibility)","ui":"settings_lora"}, - {"id":"","label":"LoRA native apply to text encoder","localized":"","hint":"","ui":"settings_lora"}, {"id":"","label":"LoRA native fuse with model","localized":"","hint":"Merge LoRA into the model for lower memory usage.

Warning: After removing or switching a LoRA, you may still see its style in generated images. To get a clean model, reload it from the model selector.","ui":"settings_lora"}, {"id":"","label":"LoRA diffusers fuse with model","localized":"","hint":"Merge LoRA into the model for lower memory usage and torch.compile compatibility.

Warning: After removing or switching a LoRA, you may still see its style in generated images. To get a clean model, reload it from the model selector.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized apply method","localized":"","hint":"How networks are applied to SDNQ-quantized model weights:
- exact: adapters are carried alongside the quantized weights at full precision; apply and removal are exact and the quantized weights are never modified. The carried factors take additional VRAM, growing with adapter rank, size and count
- requantize: adapters are merged into the quantized weights, matching the behavior of earlier releases. Uses no additional VRAM (a weight backup for network removal is held in system RAM); on models quantized below 8 bits rounding typically loses much of the adapter effect, with strong adapters retaining more

With requantize selected, the host rank, calibration and cache options below have no effect.

Default is exact.","ui":"settings_lora"}, From cf36f879a1eac60a80965a5c5b303f72d1504762 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 06:09:13 +0100 Subject: [PATCH 07/23] fix(lora): harden select stack modes on quantized and offloaded models Select pairs now ride the svd side channel on any SDNQ linear, not only sub-8-bit ones: quantized backups are packed tensors, so the weight rewrite path cannot recompute a winner from them and left layers stripped mid-requantize. The sub-8-bit gate stays for dense hosting, where requantize retains the delta at 8 bits and above. Weight selection now only serves unquantized modules: finalize iterates a snapshot so dead entries drop cleanly, materializes balanced-offload modules before rewriting weights and skips modules with stripped or quantized weights instead of corrupting the layer. --- modules/lora/lora_sdnq.py | 14 ++++++++++---- modules/lora/lora_stack.py | 20 ++++++++++++++++++-- modules/lora/networks.py | 2 +- test/test-sdnq-lora-factors.py | 34 +++++++++++++++++++++++++++++++++- 4 files changed, 62 insertions(+), 8 deletions(-) diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index df2e7d2b0..b58b26799 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -258,8 +258,8 @@ def append_factors(self, ups, downs): return segments, deq.use_quantized_matmul -def host_candidate(self, network_layer_name, wanted_names): - """True when a non-factorable set on this layer should be hosted as a truncated svd.""" +def select_candidate(self, network_layer_name, wanted_names): + """True when this layer can carry a set on the svd channel; select pairs ride it at any bit width.""" if not enabled(): return False if int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0) <= 0: @@ -268,11 +268,17 @@ def host_candidate(self, network_layer_name, wanted_names): return False if wanted_names == (): return False + return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks) + + +def host_candidate(self, network_layer_name, wanted_names): + """True when a non-factorable set on this layer should be hosted as a truncated svd.""" + if not select_candidate(self, network_layer_name, wanted_names): + return False from sdnq.common import dtype_dict if dtype_dict[self.sdnq_dequantizer.weights_dtype]['num_bits'] >= 8: return False # requantize retains most of the delta at 8 bits and above; truncation would lose more than it saves - - return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks) + return True def apply_cached(self, network_layer_name, wanted_names): diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index a9899b1d4..b74b25fde 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -248,12 +248,21 @@ def layer_flip_step(scores, total_steps): return total_steps +def materialize_model(): + """Weight-kind selection rewrites module weights outside the activation walk; rebuild balanced-offload modules real first (mirrors network_activate).""" + from modules import sd_models + if getattr(shared.opts, 'diffusers_offload_mode', None) == 'balanced' and getattr(shared, 'sd_model', None) is not None: + sd_models.apply_balanced_offload(shared.sd_model, force=True) + + def finalize(total_steps): """Build the inverted flip map for the pass; select-mode layers start at their step-0 winner.""" state['total_steps'] = int(total_steps) state['gamma'] = (state['gamma_num'] / state['gamma_den']) if state['gamma_den'] > 0 else 1.0 state['flips'] = {} - for layer_name, entry in state['entries'].items(): + if any(e['kind'] == 'weight' for e in state['entries'].values()): + materialize_model() + for layer_name, entry in list(state['entries'].items()): # snapshot: apply_selection drops entries whose module died flip_at = layer_flip_step(entry['scores'], state['total_steps']) initial = 1 if flip_at == 0 else 0 apply_selection(layer_name, entry, initial) @@ -298,6 +307,9 @@ def apply_selection(layer_name, entry, winner): def weight_selection(module, entry, winner): from modules.lora import lora_common as l from modules.lora.lora_apply import network_apply_weights + if getattr(module, 'sdnq_dequantizer', None) is not None: + warn_once('select-sdnq-weight', 'Network stack: flip=skipped layer=quantized') # quantized backups are packed tensors; only the segment path can flip them + return backup = getattr(module, 'network_weights_backup', None) if not isinstance(backup, torch.Tensor): # fuse mode keeps a bool sentinel, not a pristine copy warn_once('select-nobackup', 'Network stack: flip=skipped backup=none') @@ -306,6 +318,10 @@ def weight_selection(module, entry, winner): net_module = net.modules.get(entry['layer'], None) if net is not None else None if net_module is None: return - device = module.weight.device + weight = getattr(module, 'weight', None) + if weight is None or weight.is_meta: + warn_once('select-offloaded', 'Network stack: flip=skipped weight=offloaded') + return + device = weight.device updown = net_module.calc_updown(backup.to(device))[0] network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 6e442f5f3..818f3c100 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -112,7 +112,7 @@ def network_activate(include=None, exclude=None): device = group_offload_strip(sd_model, component, group_stripped) calced = False # tracks whether this iteration assembled the delta, so the fallthrough reuses it instead of recomputing if select_active and component_wanted and not network_layer_name.startswith('lora_te'): - if lora_sdnq.host_candidate(module, network_layer_name, component_wanted): # sub-8-bit SDNQ pairs ride the channel as separate segments + if lora_sdnq.select_candidate(module, network_layer_name, component_wanted): # SDNQ pairs ride the channel as separate segments at any bit width; weight rewrites cannot flip a quantized layer weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is not None and not isinstance(weights_backup, bool): network_apply_weights(module, None, None, device=device) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index d48e660c1..a94e38b68 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -1749,6 +1749,38 @@ def test_select_gated_off_when_compiled(): return True +def test_select_finalize_drops_dead_module(): + import weakref + layer = build_layer('uint4') + n1, n2, _D1, _D2 = select_pair(layer, seed0=61, seed1=62) + with mock_model(lin=layer), select_mode('klora'): + activate(n1, n2) + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None, 'pair must register before the module dies' + entry['module'] = weakref.ref(torch.nn.Linear(2, 2)) # referent dies immediately: simulates offload re-wraps replacing a registered module + assert entry['module']() is None + lora_stack.reset(12) + assert 'lora_transformer_test' not in lora_stack.state['entries'], 'a dead module must drop its entry without breaking finalize' + activate() + return True + + +def test_select_int8_pair_rides_segments(): + layer = build_layer('int8') + n1, n2, D1, D2 = select_pair(layer, seed0=63, seed1=64) + with mock_model(lin=layer), select_mode('klora'): + Wdq0 = dq(layer) + activate(n1, n2) + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None and entry['kind'] == 'factor', 'an int8 pair must ride svd segments, not weight rewrites' + lora_stack.reset(16) + eff = dq(layer) - Wdq0 + assert max(rho_of(eff, D1), rho_of(eff, D2)) > 0.99, 'initial selection must deliver one exact per-net delta' + activate() + assert torch.equal(dq(layer), Wdq0), 'removal must restore bit-exact' + return True + + def test_select_gate_dormant_without_pair(): with select_mode('klora'): assert lora_stack.select_possible(1) is False, 'a single network must leave the fuse gate alone' @@ -2009,7 +2041,7 @@ def run_tests(): for fn in [test_select_flip_schedule_end_to_end, test_select_initial_style_when_ramp_starts_won, test_select_flip_is_inplace_and_shape_stable, test_select_matmul_transposed_layout, test_select_per_net_hosted_pair, test_select_reset_restores_initial_state, test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled, - test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, + test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===') From 30ca66fe5fc715ec726cef7db30af94fa3e3164f Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 19 Jul 2026 17:02:38 +0100 Subject: [PATCH 08/23] fix(lora): host dense stack deltas on sdnq at any bit width Requantizing a dense-combined delta into 8-bit weights is checkpoint-fragile: on some checkpoints the round trip visibly damages the render while the same combination hosted on the svd channel is clean. Dense-mode sets with two or more contributing networks on a layer now ride the hosted path regardless of bit width; single-set behavior at 8 bits and above is unchanged. - host_candidate: dense multi-net layers qualify at any width - suite: dense pair at int8 hosts; single non-factorable set at int8 keeps the requantize fallback --- modules/lora/lora_sdnq.py | 13 +++++---- test/test-sdnq-lora-factors.py | 51 ++++++++++++++++++++++++++++++++-- 2 files changed, 57 insertions(+), 7 deletions(-) diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index b58b26799..45a514693 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -272,12 +272,15 @@ def select_candidate(self, network_layer_name, wanted_names): def host_candidate(self, network_layer_name, wanted_names): - """True when a non-factorable set on this layer should be hosted as a truncated svd.""" + """True when this layer's set should ride the svd channel as a truncated svd: non-factorable sets below 8 bits, dense-combined sets at any width.""" if not select_candidate(self, network_layer_name, wanted_names): return False + if lora_stack.mode() in lora_stack.DENSE_MODES and not network_layer_name.startswith('lora_te'): + if sum(1 for net in l.loaded_networks if net.modules.get(network_layer_name, None) is not None) >= 2: + return True # combined deltas host at any width: requantizing them is checkpoint-fragile, while single-adapter requantize is well retained from sdnq.common import dtype_dict if dtype_dict[self.sdnq_dequantizer.weights_dtype]['num_bits'] >= 8: - return False # requantize retains most of the delta at 8 bits and above; truncation would lose more than it saves + return False # requantize retains most of a single set's delta at 8 bits and above; truncation would lose more than it saves return True @@ -312,7 +315,7 @@ def apply_cached(self, network_layer_name, wanted_names): factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) if factors is not None: members.append(factors) - if len(members) == 0 and self.svd_up is None: + if not stack_dense and len(members) == 0 and self.svd_up is None: step = float(self.scale.detach().float().mean()) if step > 0 and rms / step > REQUANT_RATIO and energy < REQUANT_ENERGY: return None # routed to the grid: the caller assembles the delta and requantizes @@ -369,9 +372,9 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): # cut: both terms must agree, since a thin delta rounds away on the grid however # low its capture, and a low-rank delta hosts exactly however fat it is. Scoped to # sets the side-channel would otherwise carry whole: factorable members ride - # exactly. + # exactly and dense-combined deltas stay hosted at any magnitude. delta_rms = float(updown.detach().float().square().mean().sqrt()) - maybe_requant = len(members) == 0 and self.svd_up is None + maybe_requant = not stack_dense and len(members) == 0 and self.svd_up is None if maybe_requant: step = float(self.scale.detach().float().mean()) maybe_requant = step > 0 and delta_rms / step > REQUANT_RATIO diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index a94e38b68..8dcae106c 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -872,6 +872,19 @@ def test_route_svd_checkpoint_keeps_hosting(): return True +def test_route_dense_stack_keeps_hosting(): + layer = build_layer('uint4') + torch.manual_seed(27) + D1 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + D2 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + net1, net2 = make_dense_net('df1', layer, D1), make_dense_net('df2', layer, D2) + with host_rank(256), stack_mode('ties'), mock_model(lin=layer): + activate(net1, net2) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'dense-combined deltas host at any magnitude' + activate() + return True + + def test_route_replay_from_cache(): import tempfile layer = build_layer('uint4') @@ -1515,6 +1528,39 @@ def test_dense_two_plain_loras_hosted_not_summed(): return True +def test_dense_pair_hosts_at_int8(): + layer = build_layer('int8') + A1, B1, D1 = make_delta(seed=41, sigma=1e-2) + A2, B2, D2 = make_delta(seed=42, sigma=1e-2) + n1 = make_net('ti1', layer, A1, B1) + n2 = make_net('ti2', layer, A2, B2) + with host_rank(64), mock_model(lin=layer): + Wdq0 = dq(layer) + with stack_mode('ties', dens=0.5): + activate(n1, n2) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'dense pair at int8 must host, not requantize' + assert getattr(layer, 'network_weights_backup', None) is None, 'hosted dense pair must not take a weight backup' + eff = dq(layer) - Wdq0 + activate() + assert torch.equal(dq(layer), Wdq0), 'removal must restore bit-exact' + with stack_mode('ties', dens=0.5): + ref = lora_stack.combine([('ti1', D1), ('ti2', D2)], 'lora_transformer_test') + assert rho_of(eff, ref) > 0.8, f'hosted int8 ties delta must track the ties reference, rho={rho_of(eff, ref):.3f}' + return True + + +def test_dense_single_nonfactorable_int8_keeps_requantize(): + layer = build_layer('int8') + _A, _B, D = make_delta(sigma=3e-3) + net = make_dense_net('ti8solo', layer, D) + with host_rank(256), mock_model(lin=layer), stack_mode('ties', dens=0.5): + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'a single non-factorable set at int8 must keep the requantize path even under a dense mode' + assert isinstance(getattr(layer, 'network_weights_backup', None), torch.Tensor), 'the requantize fallback must take the backup' + activate() + return True + + def test_single_net_ignores_dense_mode(): layer = build_layer('uint4') A, B, D = make_delta(seed=33) @@ -2018,7 +2064,7 @@ def run_tests(): for fn in [test_hosted_low_rank_delta_is_kept, test_hosted_dense_delta_beats_requant, test_hosted_skips_int8, test_hosted_disabled_by_option, test_hosted_transitions_and_rng_isolation, test_route_fat_dense_delta_requantizes, test_route_rule_terms_gate_both_ways, test_route_low_rank_fat_delta_stays_hosted, - test_route_mixed_set_keeps_hosting, test_route_svd_checkpoint_keeps_hosting, + test_route_mixed_set_keeps_hosting, test_route_svd_checkpoint_keeps_hosting, test_route_dense_stack_keeps_hosting, test_route_replay_from_cache, test_hosted_null_tail_collapses_to_effective_rank, test_hosted_flat_spectrum_keeps_cap]: run_test(CAT_HOST, fn) log.warning('=== Calibration ===') @@ -2034,7 +2080,8 @@ def run_tests(): run_test(CAT_FCACHE, fn) log.warning('=== Stack modes: dense ===') for fn in [test_ties_sign_consensus_drops_conflicts, test_dare_mask_is_deterministic_across_calls, test_dare_rescales_by_inverse_density, - test_magnitude_prune_keeps_top_density, test_dense_two_plain_loras_hosted_not_summed, test_single_net_ignores_dense_mode, + test_magnitude_prune_keeps_top_density, test_dense_two_plain_loras_hosted_not_summed, + test_dense_pair_hosts_at_int8, test_dense_single_nonfactorable_int8_keeps_requantize, test_single_net_ignores_dense_mode, test_te_layer_stays_plain_sum, test_sum_mode_keeps_exact_stacking]: run_test(CAT_STACK, fn) log.warning('=== Stack modes: select ===') From b90afe1253e56efbf4625960f12a83595c35c1f6 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 19 Jul 2026 18:04:15 +0100 Subject: [PATCH 09/23] fix(lora): report the effective weight-state mode in load logs The mode field printed the configured fuse-or-backup strategy, which predates the factor path and reads as mode=backup on loads that took no backup at all. It now reports what the load actually holds: backup when weight backups were taken, fuse when fusing is active, factor when the whole load rode the svd channel and unload just drops factors. --- modules/lora/extra_networks_lora.py | 4 ++-- modules/lora/networks.py | 10 ++++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index b08dbfaf8..7d277e413 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -259,7 +259,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): if has_changed: jobid = shared.state.begin('LoRA') if len(l.previously_loaded_networks) > 0: - log.info(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if lora_overrides.fuse_native() else "backup"}') + log.info(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} mode={networks.effective_mode()}') networks.network_deactivate(include, exclude) networks.network_activate(include, exclude) debug_log(f'Network change: type=LoRA previous={[n.name for n in l.previously_loaded_networks]} current={[n.name for n in l.loaded_networks]}') @@ -272,7 +272,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): prompt(p) if has_changed and len(include) == 0: # print only once actual_method = 'native' if any(len(n.modules) > 0 for n in l.loaded_networks) else load_method - log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={"fuse" if lora_overrides.fuse_native() else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"') + log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={networks.effective_mode()} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"') def deactivate(self, p, force=False): if len(lora_diffusers.diffuser_loaded) > 0 and (shared.opts.lora_force_reload or force): diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 818f3c100..c0ffef6b1 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -222,6 +222,7 @@ def network_activate(include=None, exclude=None): lora_sdnq.report_fallbacks() native_active = len(l.loaded_networks) > 0 refused_writes = refused + l.last_backup_size = backup_size l.timer.activate += time.time() - t0 if refused > 0: log.error(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} weights={applied_weight} bias={applied_bias} refused={refused} network partially applied') @@ -232,6 +233,15 @@ def network_activate(include=None, exclude=None): sd_models.set_diffuser_offload(sd_model, op="model") +def effective_mode(): + """Weight-state label for load logs: backup and fuse say how touched weights restore, factor means the whole load rode the svd channel and unload just drops factors.""" + if getattr(l, 'last_backup_size', 0) > 0: + return 'backup' + if lora_overrides.fuse_native(): + return 'fuse' + return 'factor' + + def network_deactivate(include=None, exclude=None): if exclude is None: exclude = [] From c75e9410fe76f99379b6659e19e0cbf65d587306 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 19 Jul 2026 19:10:20 +0100 Subject: [PATCH 10/23] fix(lora): silence offload re-init logging on network changes Rebuilding balanced offload before touching weights reconstructs the OffloadHook, whose constructor prints the op=init banner and module inventory meant for model load, so every network switch replayed the full load-time announcement. The hook constructor and the model summary now honor the silent flag and the network activate, deactivate, and selection paths pass it; real model loads keep the full output. --- modules/lora/lora_stack.py | 2 +- modules/lora/networks.py | 4 ++-- modules/sd_offload_balanced.py | 11 ++++++----- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index b74b25fde..9b7d22182 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -252,7 +252,7 @@ def materialize_model(): """Weight-kind selection rewrites module weights outside the activation walk; rebuild balanced-offload modules real first (mirrors network_activate).""" from modules import sd_models if getattr(shared.opts, 'diffusers_offload_mode', None) == 'balanced' and getattr(shared, 'sd_model', None) is not None: - sd_models.apply_balanced_offload(shared.sd_model, force=True) + sd_models.apply_balanced_offload(shared.sd_model, force=True, silent=True) def finalize(total_steps): diff --git a/modules/lora/networks.py b/modules/lora/networks.py index c0ffef6b1..a1af90518 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -63,7 +63,7 @@ def network_activate(include=None, exclude=None): sd_models.disable_offload(sd_model) sd_models.move_model(sd_model, device=devices.cpu) elif shared.opts.diffusers_offload_mode == "balanced": - sd_model = sd_models.apply_balanced_offload(sd_model, force=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights + sd_model = sd_models.apply_balanced_offload(sd_model, force=True, silent=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights group_offload = shared.opts.diffusers_offload_mode == "group" group_stripped = {} device = None @@ -259,7 +259,7 @@ def network_deactivate(include=None, exclude=None): sd_models.disable_offload(sd_model) sd_models.move_model(sd_model, device=devices.cpu) elif shared.opts.diffusers_offload_mode == "balanced": - sd_model = sd_models.apply_balanced_offload(sd_model, force=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights + sd_model = sd_models.apply_balanced_offload(sd_model, force=True, silent=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights group_offload = shared.opts.diffusers_offload_mode == "group" group_stripped = {} modules = {} diff --git a/modules/sd_offload_balanced.py b/modules/sd_offload_balanced.py index 589ff3522..23a7017c6 100644 --- a/modules/sd_offload_balanced.py +++ b/modules/sd_offload_balanced.py @@ -14,7 +14,7 @@ import modules.sd_offload_state as s class OffloadHook(accelerate.hooks.ModelHook): - def __init__(self, checkpoint_name): + def __init__(self, checkpoint_name, silent=False): if shared.opts.diffusers_offload_max_gpu_memory > 1: shared.opts.diffusers_offload_max_gpu_memory = 0.75 if shared.opts.diffusers_offload_max_cpu_memory > 1: @@ -32,8 +32,9 @@ class OffloadHook(accelerate.hooks.ModelHook): self.last_pre = None self.last_post = None self.last_cls = None - gpu = f'{(shared.gpu_memory * shared.opts.diffusers_offload_min_gpu_memory):.2f}-{(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory):.2f}:{shared.gpu_memory:.2f}' - log.info(f'Offload: type=balanced op=init watermark={self.min_watermark}-{self.max_watermark} gpu={gpu} cpu={shared.cpu_memory:.3f} limit={shared.opts.cuda_mem_fraction:.2f} always={self.offload_always} never={self.offload_never} pre={shared.opts.diffusers_offload_pre} streams={shared.opts.diffusers_offload_streams}') + if not silent: + gpu = f'{(shared.gpu_memory * shared.opts.diffusers_offload_min_gpu_memory):.2f}-{(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory):.2f}:{shared.gpu_memory:.2f}' + log.info(f'Offload: type=balanced op=init watermark={self.min_watermark}-{self.max_watermark} gpu={gpu} cpu={shared.cpu_memory:.3f} limit={shared.opts.cuda_mem_fraction:.2f} always={self.offload_always} never={self.offload_never} pre={shared.opts.diffusers_offload_pre} streams={shared.opts.diffusers_offload_streams}') self.validate() super().__init__() @@ -255,7 +256,7 @@ def apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, forc checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else sd_model.__class__.__name__ if force or (s.offload_hook_instance is None) or (s.offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory) or (s.offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory) or (checkpoint_name != s.offload_hook_instance.checkpoint_name): cached = False - s.offload_hook_instance = OffloadHook(checkpoint_name) + s.offload_hook_instance = OffloadHook(checkpoint_name, silent=silent) if cached and shared.opts.diffusers_offload_pre: s.debug_move('Offload: type=balanced op=apply skip') @@ -277,6 +278,6 @@ def apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, forc process_timer.add('offload', t1 - t0) fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access s.debug_move(f'Apply offload: time={t1 - t0:.2f} type=balanced fn={fn}') - if not cached: + if not cached and not silent: log.info(f'Model class={sd_model.__class__.__name__} modules={len(s.offload_hook_instance.offload_map)} size={s.offload_hook_instance.model_size():.3f}') return sd_model From f156026150a9f0cfb795de91edf9e76d6850fd46 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 20 Jul 2026 21:57:01 +0100 Subject: [PATCH 11/23] feat(lora): balance estlora layer scores by network magnitude EST-LoRA scores each layer by squared Frobenius energy, so a magnitude gap between the two networks enters squared and the louder network wins nearly every layer, starving the quieter one. The style side is now scaled by the total-energy ratio (mirroring klora's gamma), making selection scale-invariant so a network cannot take layers on magnitude alone. On the krea2 subject+style pair this lifts the style network from 18% to 65% of the layer-step budget. - lora_stack: accumulate per-mode energy totals, apply the balance in the est ramp - test: content-louder est pair now hands over mid-schedule where raw scoring never would - locale: note the est magnitude balance, and that a select mode gives each layer to one network so both can be under-applied, while dense modes blend more fully --- modules/lora/lora_stack.py | 13 ++++++++++-- test/test-sdnq-lora-factors.py | 39 ++++++++++++++++++++++++++++++---- ui/locale/locale_en.json | 4 ++-- 3 files changed, 48 insertions(+), 8 deletions(-) diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index 9b7d22182..55453c858 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -30,7 +30,7 @@ KLORA_BETA = 0.5 # the paper's fixed ramp offset; only the slope is user-tunable ROW_CHUNK = 512 # fp32 interiors run in first-dim slices; also fixes the DARE draw sequence SAMPLE_CAP = 1 << 22 # strided subsample bound for magnitude quantiles (full-size quantile exceeds torch limits) -state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_num': 0.0, 'gamma_den': 0.0, 'total_steps': 0, 'finalized': False} +state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_num': 0.0, 'gamma_den': 0.0, 'gamma_e': 1.0, 'gamma_e_num': 0.0, 'gamma_e_den': 0.0, 'total_steps': 0, 'finalized': False} warned: set = set() @@ -203,6 +203,9 @@ def clear(): state['gamma'] = 1.0 state['gamma_num'] = 0.0 state['gamma_den'] = 0.0 + state['gamma_e'] = 1.0 + state['gamma_e_num'] = 0.0 + state['gamma_e_den'] = 0.0 state['total_steps'] = 0 state['finalized'] = False @@ -223,6 +226,9 @@ def register(layer_name, module, kind, scores, segments=None, nets=None, abs_sum if abs_sums is not None: state['gamma_num'] += abs_sums[0] state['gamma_den'] += abs_sums[1] + if mode() == 'estlora': # est scores ARE the per-layer energies; their totals give the scale-invariant balance + state['gamma_e_num'] += scores[0] + state['gamma_e_den'] += scores[1] state['entries'][layer_name] = entry state['finalized'] = False @@ -242,8 +248,10 @@ def layer_flip_step(scores, total_steps): if ss * ramp > sc: return step else: # estlora: content keeps the layer while sc >= gamma_t * ss + # est energies are ||dW||^2, so a magnitude gap enters squared; balance the style side by + # the total-energy ratio (mirrors klora's gamma) so the louder adapter cannot win on scale alone ramp = ramp_alpha() * t + (1.0 - manual_discrepancy()) - if sc < ramp * ss: + if sc < ramp * ss * state['gamma_e']: return step return total_steps @@ -259,6 +267,7 @@ def finalize(total_steps): """Build the inverted flip map for the pass; select-mode layers start at their step-0 winner.""" state['total_steps'] = int(total_steps) state['gamma'] = (state['gamma_num'] / state['gamma_den']) if state['gamma_den'] > 0 else 1.0 + state['gamma_e'] = (state['gamma_e_num'] / state['gamma_e_den']) if state['gamma_e_den'] > 0 else 1.0 state['flips'] = {} if any(e['kind'] == 'weight' for e in state['entries'].values()): materialize_model() diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 8dcae106c..b08f2988b 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -1669,14 +1669,45 @@ def test_select_flip_schedule_end_to_end(): def test_select_initial_style_when_ramp_starts_won(): + # scale-invariant selection means no single isolated layer starts style-won on magnitude alone + # (that is the balance working), so force flip_step 0 directly and assert the initial selection honors it layer = build_layer('uint4') - n1, n2, _D1, D2 = select_pair(seed0=43, seed1=44, scale1=8.0, layer=layer) # style delta dominates + n1, n2, _D1, D2 = select_pair(seed0=43, seed1=44, layer=layer) with mock_model(lin=layer), select_mode('estlora', alpha=1.0, disc=0.5): Wdq0 = dq(layer) activate(n1, n2) - lora_stack.reset(20) + orig = lora_stack.layer_flip_step + lora_stack.layer_flip_step = lambda scores, total: 0 # this layer's crossover is step 0 + try: + lora_stack.reset(20) + finally: + lora_stack.layer_flip_step = orig eff = dq(layer) - Wdq0 - assert rho_of(eff, D2) > 0.99, 'a layer whose style side wins at step 0 must start style-selected' + assert rho_of(eff, D2) > 0.99, 'a layer whose flip step is 0 must start style-selected' + return True + + +def test_estlora_energy_balance_defeats_magnitude(): + layer = build_layer('uint4') + n1, n2, _D1, D2 = select_pair(layer, seed0=51, seed1=52, scale1=0.33) # content ~3x louder than style + with mock_model(lin=layer), select_mode('estlora', alpha=1.5, disc=0.5): + Wdq0 = dq(layer) + activate(n1, n2) + lora_stack.reset(20) + entry = lora_stack.state['entries']['lora_transformer_test'] + assert lora_stack.state['gamma_e'] > 1.5, f'content-louder pair must give gamma_e>1: {lora_stack.state["gamma_e"]:.2f}' + balanced = lora_stack.layer_flip_step(entry['scores'], 20) + saved = lora_stack.state['gamma_e'] + lora_stack.state['gamma_e'] = 1.0 # paper-faithful est: energies compared raw + raw = lora_stack.layer_flip_step(entry['scores'], 20) + lora_stack.state['gamma_e'] = saved + assert raw == 20, f'without balance the squared magnitude gap keeps content the whole schedule, got {raw}' + assert 0 < balanced < 20, f'the energy balance must let the quieter style win mid-schedule, got {balanced}' + for s in range(20): + lora_stack.on_step(s) + assert rho_of(dq(layer) - Wdq0, D2) > 0.99, 'after the balanced flip the style delta must be selected' + activate() + assert torch.equal(dq(layer), Wdq0) return True @@ -2085,7 +2116,7 @@ def run_tests(): test_te_layer_stays_plain_sum, test_sum_mode_keeps_exact_stacking]: run_test(CAT_STACK, fn) log.warning('=== Stack modes: select ===') - for fn in [test_select_flip_schedule_end_to_end, test_select_initial_style_when_ramp_starts_won, test_select_flip_is_inplace_and_shape_stable, + for fn in [test_select_flip_schedule_end_to_end, test_select_initial_style_when_ramp_starts_won, test_estlora_energy_balance_defeats_magnitude, test_select_flip_is_inplace_and_shape_stable, test_select_matmul_transposed_layout, test_select_per_net_hosted_pair, test_select_reset_restores_initial_state, test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled, test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 2dd87c144..f52b58a91 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -867,10 +867,10 @@ {"id":"","label":"LoRA quantized host rank","localized":"","hint":"Maximum rank used to carry adapter types that are not natively low-rank (LoKR, LoHA, OFT, DoRA) alongside the quantized weights instead of merging them in.
Higher values retain more of the adapter at proportionally more memory. Plain LoRA files are carried exactly at their own rank.

Applies only to SDNQ models quantized below 8 bits, where merging erases most of the adapter; at 8 bits and above merging retains it and hosting is skipped.

0 disables hosting and merges every adapter into the quantized weights.

Default is 256.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized host calibration","localized":"","hint":"Collects per-channel activation statistics from the model's own generations and uses them to focus hosted-adapter truncation on the channels with the strongest activations.
Statistics accumulate in the background on models quantized below 8 bits, persist per checkpoint, and raise delivered adapter fidelity at the same LoRA quantized host rank, most at low ranks.

Capture is skipped while the model is compiled; previously cached statistics still apply.

Enabled by default.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized host cache","localized":"","hint":"Disk space in GB for caching computed hosting factors.
A cached set skips the truncation math on the next load; least recently used entries are evicted once the budget is exceeded.

0 disables the cache.

Default is 10.","ui":"settings_lora"}, - {"id":"","label":"LoRA stack mode","localized":"","hint":"How multiple networks targeting the same layer are combined:
- sum: adds all contributions
- ties: keeps each network's strongest elements and merges only where signs agree
- dare_ties: randomly drops elements, rescales the survivors, then merges where signs agree
- dare_linear: randomly drops elements, rescales the survivors and sums
- magnitude_prune: keeps each network's strongest elements and sums
- klora / estlora: assign each layer to one of exactly two networks, the first in the prompt as subject and the second as style, shifting from subject toward style over the sampling steps

Kept fractions are set by LoRA stack density; the subject-to-style shift by LoRA stack ramp and LoRA stack discrepancy.

Applies to the native load path; other load methods and text encoder networks always combine as sum. Selection modes fall back to sum unless exactly two networks are loaded, or when model compile is active.

Default is sum.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack mode","localized":"","hint":"How multiple networks targeting the same layer are combined:
- sum: adds all contributions
- ties: keeps each network's strongest elements and merges only where signs agree
- dare_ties: randomly drops elements, rescales the survivors, then merges where signs agree
- dare_linear: randomly drops elements, rescales the survivors and sums
- magnitude_prune: keeps each network's strongest elements and sums
- klora / estlora: assign each layer to one of exactly two networks, the first in the prompt as subject and the second as style, shifting from subject toward style over the sampling steps

Each layer is given to a single network at a time, so a subject and a style that both need sustained strength can end up under-applied. For reliable blending of two strong networks, sum, ties and dare_ties apply every network throughout and combine more fully.

Kept fractions are set by LoRA stack density; the subject-to-style shift by LoRA stack ramp and LoRA stack discrepancy.

Applies to the native load path; other load methods and text encoder networks always combine as sum. Selection modes fall back to sum unless exactly two networks are loaded, or when model compile is active.

Default is sum.","ui":"settings_lora"}, {"id":"","label":"LoRA stack density","localized":"","hint":"Fraction of elements each network keeps under the ties, dare_ties, dare_linear and magnitude_prune stack modes.
Lower values keep only the strongest contributions and reduce interference between networks at the cost of per-network detail. The dare variants drop at random and rescale the survivors to preserve expected strength.

Default is 0.5.","ui":"settings_lora"}, {"id":"","label":"LoRA stack ramp","localized":"","hint":"Slope of the subject-to-style shift across the sampling steps in the klora and estlora stack modes.
Higher values shift layers to the style network earlier and more broadly; lower values keep the subject network dominant for longer.

0 keeps the balance fixed for the whole generation.

Default is 1.5.","ui":"settings_lora"}, - {"id":"","label":"LoRA stack discrepancy","localized":"","hint":"Stand-in for the measured style separation the estlora stack mode would otherwise derive from data.
Higher values keep layers with the subject network longer; lower values let the style network take layers earlier.

Applies only when LoRA stack mode is estlora.

Default is 0.5.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack discrepancy","localized":"","hint":"Stand-in for the measured style separation the estlora stack mode would otherwise derive from data.
Higher values keep layers with the subject network longer; lower values let the style network take layers earlier.

Layer scores are balanced by each network's overall strength, so a louder network does not take layers on magnitude alone.

Applies only when LoRA stack mode is estlora.

Default is 0.5.","ui":"settings_lora"}, {"id":"","label":"LoRA auto-apply tags","localized":"","hint":"Automatically add trigger words/tags from LoRA metadata to your prompt.
Set to the number of tags to auto-apply, e.g., 3 = add top 3 trigger tags.
Set to 0 to disable, -1 to add all available tags.","ui":"settings_lora"}, {"id":"","label":"LoRA memory cache","localized":"","hint":"How many LoRAs to keep in network for future use before requiring reloading from storage","ui":"settings_lora"}, {"id":"","label":"LoRA add hash info to metadata","localized":"","hint":"Include LoRA file hashes in generated image metadata.
Useful for reproducibility and tracking which exact LoRA versions were used.","ui":"settings_lora"}, From 3218740b20abc9538ce0a4654640972523c2db72 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 20 Jul 2026 23:33:03 +0100 Subject: [PATCH 12/23] fix(lora): surface the stack mode as an explicit field in load logs The active stack mode was only visible inside the stack= token of the trace-level network check line. Add it to the load summary as its own stack= field alongside method, mode, te and unet, carrying the mode and its tuning (ties:0.50, klora:1.50:0.50, sum). Non-native loads report sum, since those paths always combine as sum regardless of the setting. --- modules/lora/extra_networks_lora.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index 7d277e413..c0fce6fb4 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -272,7 +272,8 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): prompt(p) if has_changed and len(include) == 0: # print only once actual_method = 'native' if any(len(n.modules) > 0 for n in l.loaded_networks) else load_method - log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={networks.effective_mode()} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"') + stack = lora_stack.signature() if actual_method == 'native' else 'sum' # non-native paths always combine as sum + log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={networks.effective_mode()} stack={stack} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"') def deactivate(self, p, force=False): if len(lora_diffusers.diffuser_loaded) > 0 and (shared.opts.lora_force_reload or force): From 82e3c1d20f7e5557a3cbe03bf0c67f0cf7895f78 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 25 Jul 2026 22:14:47 +0100 Subject: [PATCH 13/23] feat(lora): log select stack schedules Select modes left no trace distinguishable from plain summation: apply_select counted its layers on the exact path, and the mode field in the load summary reflects the requested setting rather than what executed. A flip count can only come from a populated schedule. - report layers, initial style picks, flips, steps and gamma from finalize - deduplicate on content, since the schedule rebuilds on every pass - give select its own apply counter instead of inflating apply=exact --- modules/lora/lora_sdnq.py | 6 +++++- modules/lora/lora_stack.py | 11 ++++++++++- modules/lora/networks.py | 1 + 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 45a514693..375e8d05a 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -56,6 +56,7 @@ fallback_layers: list[str] = [] hosted_layers: list[tuple[str, float, bool]] = [] hosted_ranks: list[int] = [] factor_layers: list[str] = [] +select_layers: list[str] = [] routed_layers: list[str] = [] REQUANT_RATIO = 0.30 # delta rms over mean grid step above which requantize can retain the delta @@ -507,7 +508,7 @@ def apply_select(self, network_layer_name, per_net, wanted_names): del d0, d1 segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]]) lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums) - factor_layers.append(network_layer_name) + select_layers.append(network_layer_name) # counted apart from the plain concat: both ride the svd channel but only one is a summed set return True @@ -524,6 +525,9 @@ def report_fallbacks(): if len(factor_layers) > 0: log.info(f'Network load: type=LoRA quant=sdnq apply=exact layers={len(factor_layers)}') factor_layers.clear() + if len(select_layers) > 0: + log.info(f'Network load: type=LoRA quant=sdnq apply=select layers={len(select_layers)} mode={lora_stack.mode()}') + select_layers.clear() if len(hosted_layers) > 0: energies = sorted(e for _name, e, _c in hosted_layers) median = energies[len(energies) // 2] diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index 55453c858..38f586560 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -30,7 +30,7 @@ KLORA_BETA = 0.5 # the paper's fixed ramp offset; only the slope is user-tunable ROW_CHUNK = 512 # fp32 interiors run in first-dim slices; also fixes the DARE draw sequence SAMPLE_CAP = 1 << 22 # strided subsample bound for magnitude quantiles (full-size quantile exceeds torch limits) -state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_num': 0.0, 'gamma_den': 0.0, 'gamma_e': 1.0, 'gamma_e_num': 0.0, 'gamma_e_den': 0.0, 'total_steps': 0, 'finalized': False} +state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_num': 0.0, 'gamma_den': 0.0, 'gamma_e': 1.0, 'gamma_e_num': 0.0, 'gamma_e_den': 0.0, 'total_steps': 0, 'finalized': False, 'reported': None} warned: set = set() @@ -208,6 +208,7 @@ def clear(): state['gamma_e_den'] = 0.0 state['total_steps'] = 0 state['finalized'] = False + state['reported'] = None def register(layer_name, module, kind, scores, segments=None, nets=None, abs_sums=None): @@ -271,13 +272,21 @@ def finalize(total_steps): state['flips'] = {} if any(e['kind'] == 'weight' for e in state['entries'].values()): materialize_model() + style_first = 0 for layer_name, entry in list(state['entries'].items()): # snapshot: apply_selection drops entries whose module died flip_at = layer_flip_step(entry['scores'], state['total_steps']) initial = 1 if flip_at == 0 else 0 + style_first += initial apply_selection(layer_name, entry, initial) if 0 < flip_at < state['total_steps']: state['flips'].setdefault(flip_at, []).append(layer_name) state['finalized'] = True + if len(state['entries']) > 0: # only a built schedule can carry a flip count, so this is the line that shows selection is live rather than requested + gamma = state['gamma_e'] if mode() == 'estlora' else state['gamma'] + report = (mode(), len(state['entries']), style_first, sum(len(v) for v in state['flips'].values()), state['total_steps'], round(gamma, 3)) + if report != state['reported']: # rebuilt every pass, so a batch would otherwise repeat one line per image + state['reported'] = report + log.info(f'Network load: type=LoRA stack={report[0]} layers={report[1]} style={report[2]} flips={report[3]} steps={report[4]} gamma={report[5]:.3f}') def reset(total_steps): diff --git a/modules/lora/networks.py b/modules/lora/networks.py index a1af90518..0987c679d 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -96,6 +96,7 @@ def network_activate(include=None, exclude=None): lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind lora_sdnq.hosted_layers.clear() lora_sdnq.factor_layers.clear() + lora_sdnq.select_layers.clear() backup_size = 0 for component in modules.keys(): component_wanted = wanted_names if component in components else () From 72d9fd1b7805f4c32abf373e5ae95840b25707e0 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 25 Jul 2026 23:32:59 +0100 Subject: [PATCH 14/23] fix(lora): derive select stack balances from live schedule entries The klora and estlora balance factors accumulated across registrations without ever resetting, so a multiplier change or pair swap blended the previous registration into every later schedule. Balances now sum over the live entries at finalize time, which keeps drop and re-register consistent by construction. - key flips one step early: step callbacks fire after the denoise, so the winner is now live during the crossover step forward and a final-step crossover engages instead of expiring - include the calibration toggle in the factor cache signature so a hit never replays factors computed under the other setting - fall back to summation with a warning when hosting is disabled on a quantized model instead of registering schedules that cannot flip - drop the unused score_topk helper --- modules/lora/lora_stack.py | 32 +++++--------- modules/lora/networks.py | 3 ++ test/test-sdnq-lora-factors.py | 76 +++++++++++++++++++++++++++++++++- 3 files changed, 88 insertions(+), 23 deletions(-) diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index 38f586560..ef60bc735 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -30,7 +30,7 @@ KLORA_BETA = 0.5 # the paper's fixed ramp offset; only the slope is user-tunable ROW_CHUNK = 512 # fp32 interiors run in first-dim slices; also fixes the DARE draw sequence SAMPLE_CAP = 1 << 22 # strided subsample bound for magnitude quantiles (full-size quantile exceeds torch limits) -state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_num': 0.0, 'gamma_den': 0.0, 'gamma_e': 1.0, 'gamma_e_num': 0.0, 'gamma_e_den': 0.0, 'total_steps': 0, 'finalized': False, 'reported': None} +state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_e': 1.0, 'total_steps': 0, 'finalized': False, 'reported': None} warned: set = set() @@ -183,13 +183,6 @@ def drop(layer_name): state['finalized'] = False -def score_topk(up, down, k): - """K-LoRA layer score: sum of the top-K absolute delta entries (one dense materialization).""" - d = (up.to(torch.float32) @ down.to(torch.float32)).abs().flatten() - values = torch.topk(d, min(int(k), d.numel()), sorted=False).values - return float(values.sum()), float(d.sum()) - - def score_energy(up, down): """EST layer score: squared Frobenius norm of up@down via the Gram identity, no materialization.""" u = up.to(torch.float32) @@ -201,11 +194,7 @@ def clear(): state['entries'] = {} state['flips'] = {} state['gamma'] = 1.0 - state['gamma_num'] = 0.0 - state['gamma_den'] = 0.0 state['gamma_e'] = 1.0 - state['gamma_e_num'] = 0.0 - state['gamma_e_den'] = 0.0 state['total_steps'] = 0 state['finalized'] = False state['reported'] = None @@ -219,17 +208,11 @@ def register(layer_name, module, kind, scores, segments=None, nets=None, abs_sum the two network names; the winner delta is recomputed from the layer backup at selection time. abs_sums feeds the global magnitude balance (klora gamma). """ - entry = {'layer': layer_name, 'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'nets': nets, 'stash': None} + entry = {'layer': layer_name, 'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'nets': nets, 'abs_sums': abs_sums, 'stash': None} if kind == 'factor': (s0, s1), (t0, t1), transposed = segments up = module.svd_up.data entry['stash'] = (segment_view(up, s0, s1, transposed).clone(), segment_view(up, t0, t1, transposed).clone()) - if abs_sums is not None: - state['gamma_num'] += abs_sums[0] - state['gamma_den'] += abs_sums[1] - if mode() == 'estlora': # est scores ARE the per-layer energies; their totals give the scale-invariant balance - state['gamma_e_num'] += scores[0] - state['gamma_e_den'] += scores[1] state['entries'][layer_name] = entry state['finalized'] = False @@ -267,8 +250,13 @@ def materialize_model(): def finalize(total_steps): """Build the inverted flip map for the pass; select-mode layers start at their step-0 winner.""" state['total_steps'] = int(total_steps) - state['gamma'] = (state['gamma_num'] / state['gamma_den']) if state['gamma_den'] > 0 else 1.0 - state['gamma_e'] = (state['gamma_e_num'] / state['gamma_e_den']) if state['gamma_e_den'] > 0 else 1.0 + # both balances derive from the live entries every time, so drops and re-registrations stay consistent by construction + num = sum(e['abs_sums'][0] for e in state['entries'].values() if e['abs_sums'] is not None) + den = sum(e['abs_sums'][1] for e in state['entries'].values() if e['abs_sums'] is not None) + state['gamma'] = (num / den) if den > 0 else 1.0 + e_num = sum(e['scores'][0] for e in state['entries'].values()) # est scores ARE the per-layer energies; their totals give the scale-invariant balance + e_den = sum(e['scores'][1] for e in state['entries'].values()) + state['gamma_e'] = (e_num / e_den) if e_den > 0 else 1.0 state['flips'] = {} if any(e['kind'] == 'weight' for e in state['entries'].values()): materialize_model() @@ -279,7 +267,7 @@ def finalize(total_steps): style_first += initial apply_selection(layer_name, entry, initial) if 0 < flip_at < state['total_steps']: - state['flips'].setdefault(flip_at, []).append(layer_name) + state['flips'].setdefault(flip_at - 1, []).append(layer_name) # step callbacks fire after the denoise, so the flip runs one step early to be live during the crossover step's forward state['finalized'] = True if len(state['entries']) > 0: # only a built schedule can carry a flip count, so this is the line that shows selection is live rather than requested gamma = state['gamma_e'] if mode() == 'estlora' else state['gamma'] diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 0987c679d..651fe8826 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -129,6 +129,9 @@ def network_activate(include=None, exclude=None): if task is not None: pbar.update(task, advance=1) continue + elif getattr(module, 'sdnq_dequantizer', None) is not None: # hosting disabled: quantized layers have no side-channel to carry segments and packed backups cannot flip, so the sum paths below take the layer + if any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks): + lora_stack.warn_once('select-host-disabled', f'Network stack: mode={lora_stack.mode()} quant=sdnq host=disabled fallback=sum') else: # other layers select by recomputing the winner from the pristine backup at schedule time backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse) weights_backup = getattr(module, "network_weights_backup", None) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index b08f2988b..7b4455077 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -1930,6 +1930,79 @@ def test_select_weight_kind_plain_layer(): return True +def test_select_gamma_tracks_live_entries(): + layer = build_layer('uint4') + n1, n2, _D1, _D2 = select_pair(layer, seed0=57, seed1=58) + with mock_model(lin=layer), select_mode('klora'): + activate(n1, n2) + lora_stack.reset(20) + entry = lora_stack.state['entries']['lora_transformer_test'] + live = entry['abs_sums'][0] / entry['abs_sums'][1] + assert abs(lora_stack.state['gamma'] - live) < 1e-9, f'gamma {lora_stack.state["gamma"]} vs live ratio {live}' + n2.te_multiplier = 0.5 + n2.unet_multiplier = [0.5] * 3 + activate(n1, n2) # multiplier change re-applies the pair through the drop/re-register walk + lora_stack.reset(20) + entry = lora_stack.state['entries']['lora_transformer_test'] + live2 = entry['abs_sums'][0] / entry['abs_sums'][1] + assert live2 > live * 1.5, f'halving the style multiplier must move the live ratio: {live} -> {live2}' + assert abs(lora_stack.state['gamma'] - live2) < 1e-9, f'gamma must equal the live-entry ratio, not blend with the previous registration: {lora_stack.state["gamma"]} vs {live2}' + activate() + with mock_model(lin=layer), select_mode('estlora'): + activate(n1, n2) + lora_stack.reset(20) + entry = lora_stack.state['entries']['lora_transformer_test'] + live_e = entry['scores'][0] / entry['scores'][1] + assert abs(lora_stack.state['gamma_e'] - live_e) < 1e-9 + n2.te_multiplier = 1.0 + n2.unet_multiplier = [1.0] * 3 + activate(n1, n2) + lora_stack.reset(20) + entry = lora_stack.state['entries']['lora_transformer_test'] + live_e2 = entry['scores'][0] / entry['scores'][1] + assert abs(lora_stack.state['gamma_e'] - live_e2) < 1e-9, f'energy balance must track the live registration: {lora_stack.state["gamma_e"]} vs {live_e2}' + activate() + return True + + +def test_select_host_disabled_falls_back_to_sum(): + layer = build_layer('uint4') + n1, n2, D1, D2 = select_pair(layer, seed0=53, seed1=54) + with mock_model(lin=layer), select_mode('klora'), host_rank(0): + Wdq0 = dq(layer) + activate(n1, n2) + assert not lora_stack.select_engaged(), 'hosting disabled: quantized layers cannot carry segments, nothing must schedule' + assert 'select-host-disabled' in lora_stack.warned, 'the degradation must be said once' + eff = dq(layer) - Wdq0 + assert rho_of(eff, D1 + D2) > 0.99, 'the pair must land as plain summation, not a pristine no-op' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_flip_lands_before_crossover_step(): + layer = build_layer('uint4') + n1, n2, D1, D2 = select_pair(layer, seed0=55, seed1=56) + with mock_model(lin=layer), select_mode('klora'): + Wdq0 = dq(layer) + activate(n1, n2) + total = 20 + orig = lora_stack.layer_flip_step + lora_stack.layer_flip_step = lambda scores, t: t - 1 # crossover on the final step + try: + lora_stack.reset(total) + finally: + lora_stack.layer_flip_step = orig + assert list(lora_stack.state['flips'].keys()) == [total - 2], f'end-of-step callbacks: a crossover at step k must execute at the end of step k-1, got {list(lora_stack.state["flips"].keys())}' + assert rho_of(dq(layer) - Wdq0, D1) > 0.99, 'the subject holds the layer before the flip' + for s in range(total - 1): # the callback after the second-to-last denoise is the last one that can matter + lora_stack.on_step(s) + assert rho_of(dq(layer) - Wdq0, D2) > 0.99, 'the style side must be live for the final denoise' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + CAT_COMPILE = category('compile') @@ -2120,7 +2193,8 @@ def run_tests(): test_select_matmul_transposed_layout, test_select_per_net_hosted_pair, test_select_reset_restores_initial_state, test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled, test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, - test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer]: + test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer, + test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: From 4032ff61a88cbe3071f689563437a7815b7bc4b4 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 25 Jul 2026 23:38:21 +0100 Subject: [PATCH 15/23] fix(lora): keep the select count warning quiet without networks A restore-only activation walk carries zero networks, so the networks=0 required=2 fallback warning fired on every network-free generation whenever a select stack mode was set. The gate now short-circuits at zero; the 1-and-3-network warnings that remain meaningful are unchanged. --- modules/lora/networks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 651fe8826..aa4470da3 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -91,7 +91,7 @@ def network_activate(include=None, exclude=None): with devices.inference_context(), pbar: wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () stack_sig = lora_stack.signature() + lora_sdnq.signature() # tracked beside network_current_names so settings-only stack or mechanism changes re-apply - select_active = lora_stack.active_select(len(l.loaded_networks)) + select_active = len(l.loaded_networks) > 0 and lora_stack.active_select(len(l.loaded_networks)) # restore-only walks have nothing to stack; the count warning would fire on every network-free generation applied_layers.clear() lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind lora_sdnq.hosted_layers.clear() From 84220562ee917e26d0a1a21796b754402bc48f38 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 00:07:29 +0100 Subject: [PATCH 16/23] fix(lora): count select backups once and log unridable pairs The weight-kind select branch counted its backup at its own call site and again at the shared backup call when registration fell through, so the reported backup size double-counted those layers; the size now lands with whichever path keeps the layer. A pair the svd channel cannot carry (bias delta or malformed member) previously dropped to the sum paths with no trace; the fallthrough now says so once. --- modules/lora/networks.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modules/lora/networks.py b/modules/lora/networks.py index aa4470da3..366997118 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -129,15 +129,17 @@ def network_activate(include=None, exclude=None): if task is not None: pbar.update(task, advance=1) continue + lora_stack.warn_once('select-unridable', f'Network stack: mode={lora_stack.mode()} layer="{network_layer_name}" fallback=sum') # a pair the channel cannot carry (bias delta or malformed member) sums like any unsupported set elif getattr(module, 'sdnq_dequantizer', None) is not None: # hosting disabled: quantized layers have no side-channel to carry segments and packed backups cannot flip, so the sum paths below take the layer if any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks): lora_stack.warn_once('select-host-disabled', f'Network stack: mode={lora_stack.mode()} quant=sdnq host=disabled fallback=sum') else: # other layers select by recomputing the winner from the pristine backup at schedule time - backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse) + sel_backup = network_backup_weights(module, network_layer_name, component_wanted, fuse) weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is not None and not isinstance(weights_backup, bool): per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net): + backup_size += sel_backup # counted only when this branch keeps the layer; the fallthrough re-enters the shared backup call below, which counts it then network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner applied_layers.append(network_layer_name) applied_weight += 1 From de5e94073cca152cccbc369f55a25d6b7ec12af4 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 07:24:31 +0100 Subject: [PATCH 17/23] perf(lora): chunk select scoring and cache select scores for replay Select scoring staged full fp32 copies, an abs copy and top-k workspace per layer (hundreds of MB of transients that collide with block swapping on offloaded denoisers) and recomputed scores from freshly assembled deltas on every apply, which kept select modes out of the factor-cache fast path. - score_pair: row-chunked fp32 interiors, fp64 accumulators, one device sync - select scores persist in the factor cache as additive per-layer records under the existing configuration signature - apply_select_cached and register_weight_pair_cached replay a pair without assembling deltas; the weight-kind winner is still computed at schedule time --- modules/lora/lora_factor_cache.py | 24 ++++++ modules/lora/lora_sdnq.py | 54 +++++++++++++- modules/lora/lora_stack.py | 73 ++++++++++++++---- modules/lora/networks.py | 38 ++++++---- test/test-sdnq-lora-factors.py | 119 +++++++++++++++++++++++++++++- 5 files changed, 277 insertions(+), 31 deletions(-) diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index d6c95d681..aab36c7a5 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -157,6 +157,30 @@ def fetch(network_layer_name): return entry +def lookup_scores(network_layer_name): + """Cached select scores for a layer as ((s0, s1), (a0, a1)), or None. + + Score records ride the same signature-keyed entry as factors, and the + signature already pins everything the scores depend on (pair, multipliers, + stack mode and params). No hit/miss accounting: a record saves scoring and + delta assembly, not a sketch. + """ + if state['sig'] is None: + return None + t = state['store'].get(f'{network_layer_name}.sel') + if t is None: + return None + return (float(t[0]), float(t[1])), (float(t[2]), float(t[3])) + + +def store_scores(network_layer_name, scores, abs_sums): + """Persist a select-mode score record; additive to the entry, older files upgrade on their next pass.""" + if state['sig'] is None: + return + state['store'][f'{network_layer_name}.sel'] = torch.tensor([scores[0], scores[1], abs_sums[0], abs_sums[1]], dtype=torch.float64) + state['dirty'] = True + + def store(network_layer_name, up, down, energy, calibrated, rms): """Quantize-before-use: returns the dequantized round-trip the caller must apply. diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 375e8d05a..35635b227 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -460,6 +460,54 @@ def truncate_delta(self, D, dtype): return up_h, down_h, energy, rms is not None +def apply_select_cached(self, network_layer_name, wanted_names): + """Serve a select pair from cache and live factors before the walk assembles deltas. + + A cached score record plus a factor pair per network (exact factors for + factorable members, cached truncations otherwise) rebuild the segments and + the selection registration without any ``calc_updown``. Returns None when + any piece is missing; the caller assembles and ``apply_select`` recomputes + and stores. + """ + from sdnq.quant_utils import rotate_hadamard + deq = self.sdnq_dequantizer + changed = remove_factors(self) + if wanted_names == (): + return changed + if len(l.loaded_networks) != 2: + return None + dtype = deq.result_dtype + lora_factor_cache.begin_pass(wanted_names) + rec = lora_factor_cache.lookup_scores(network_layer_name) + if rec is None: + return None + pairs, notes = [], [] + for i, net in enumerate(l.loaded_networks): + module = net.modules.get(network_layer_name, None) + if module is None: + return None + factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) + if factors is not None: + up_i, down_i = factors + if deq.use_hadamard: + down_i = rotate_hadamard(down_i.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype) + else: + cached = lora_factor_cache.lookup(f'{network_layer_name}#{i}') + if cached is None: + return None + up_i, down_i = cached[0].to(device=devices.device, dtype=dtype), cached[1].to(device=devices.device, dtype=dtype) + notes.append((f'{network_layer_name}#{i}', cached[2], cached[3])) + pairs.append((up_i, down_i)) + scores, abs_sums = rec + segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]]) + lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums) + for note in notes: + lora_factor_cache.note_hit() + hosted_layers.append(note) + select_layers.append(network_layer_name) + return True + + def apply_select(self, network_layer_name, per_net, wanted_names): """Attach two networks' contributions as separate side-channel segments for per-layer selection. @@ -502,10 +550,8 @@ def apply_select(self, network_layer_name, per_net, wanted_names): up_i, down_i = lora_factor_cache.store(key, up_i, down_i, energy, calibrated, float(D.detach().float().square().mean().sqrt())) hosted_layers.append((key, energy, calibrated)) pairs.append((up_i, down_i)) - d0 = per_net[0][1].detach().to(devices.device, torch.float32) - d1 = per_net[1][1].detach().to(devices.device, torch.float32) - scores, abs_sums = lora_stack.score_pair(d0, d1, ranks[0], ranks[1]) - del d0, d1 + scores, abs_sums = lora_stack.score_pair(per_net[0][1].detach(), per_net[1][1].detach(), ranks[0], ranks[1]) + lora_factor_cache.store_scores(network_layer_name, scores, abs_sums) segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]]) lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums) select_layers.append(network_layer_name) # counted apart from the plain concat: both ride the svd channel but only one is a summed set diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index ef60bc735..acd1117fa 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -145,20 +145,43 @@ def combine(named_deltas, layer_name): def score_pair(d0, d1, rank0, rank1): - """Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance.""" - abs_sums = (float(d0.abs().sum()), float(d1.abs().sum())) - if mode() == 'klora': - k = max(1, int(rank0) * int(rank1)) - s0 = float(torch.topk(d0.abs().flatten(), min(k, d0.numel()), sorted=False).values.sum()) - s1 = float(torch.topk(d1.abs().flatten(), min(k, d1.numel()), sorted=False).values.sum()) - else: - s0 = float(d0.float().square().sum()) - s1 = float(d1.float().square().sum()) - return (s0, s1), abs_sums + """Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance. + + Row-chunked fp32 interiors with fp64 accumulators and one device sync for + all four reductions. Full-tensor staging (fp32 copy, abs copy, top-k + workspace) peaks hundreds of MB per large layer, which collides with block + swapping on offloaded denoisers; chunking bounds the transient to the + chunk. The global top-K over per-chunk top-K candidates selects the same + element set as a whole-tensor top-K. + """ + k = max(1, int(rank0) * int(rank1)) if mode() == 'klora' else 0 + accs = [] + for d in (d0, d1): + score = torch.zeros((), device=d.device, dtype=torch.float64) + abs_sum = torch.zeros((), device=d.device, dtype=torch.float64) + cands = [] + for start in range(0, d.shape[0], ROW_CHUNK): + c = d[start:start + ROW_CHUNK].to(torch.float32).abs() # out-of-place abs: to() may alias a caller-owned fp32 tensor + abs_sum += c.sum(dtype=torch.float64) + if k: + flat = c.flatten() + cands.append(torch.topk(flat, min(k, flat.numel()), sorted=False).values) + else: + score += c.square().sum(dtype=torch.float64) + if k and cands: + allc = torch.cat(cands) if len(cands) > 1 else cands[0] + score = torch.topk(allc, min(k, allc.numel()), sorted=False).values.sum(dtype=torch.float64) + accs.append((score, abs_sum)) + packed = torch.stack([accs[0][0], accs[0][1], accs[1][0], accs[1][1]]).cpu() + return (float(packed[0]), float(packed[2])), (float(packed[1]), float(packed[3])) -def register_weight_pair(layer_name, module, per_net): - """Score and register a weight-kind selection pair; True when the layer is scheduled.""" +def register_weight_pair(layer_name, module, per_net, wanted_names=None): + """Score and register a weight-kind selection pair; True when the layer is scheduled. + + The scores persist in the factor cache when a pass identity is given, so a + later apply of the same configuration registers from the record alone. + """ from modules.lora import lora_common as l if per_net is None or len(per_net) != 2: return False @@ -172,11 +195,35 @@ def register_weight_pair(layer_name, module, per_net): return False names.append(net_name) ranks.append(int(getattr(net_module, 'dim', 0) or 0) or 64) - scores, abs_sums = score_pair(per_net[0][1].float(), per_net[1][1].float(), ranks[0], ranks[1]) + scores, abs_sums = score_pair(per_net[0][1], per_net[1][1], ranks[0], ranks[1]) + if wanted_names is not None: + from modules.lora import lora_factor_cache + lora_factor_cache.begin_pass(wanted_names) + lora_factor_cache.store_scores(layer_name, scores, abs_sums) register(layer_name, module, 'weight', scores, nets=tuple(names), abs_sums=abs_sums) return True +def register_weight_pair_cached(layer_name, module, wanted_names): + """Register a weight-kind pair from its cached score record; True when served. + + The record was stored under the same configuration signature, which pins + the loaded pair, multipliers and stack settings, so both networks are known + to target the layer and the prompt-order roles are unchanged. + """ + from modules.lora import lora_common as l + from modules.lora import lora_factor_cache + if len(l.loaded_networks) != 2: + return False + lora_factor_cache.begin_pass(wanted_names) + rec = lora_factor_cache.lookup_scores(layer_name) + if rec is None: + return False + scores, abs_sums = rec + register(layer_name, module, 'weight', scores, nets=tuple(n.name for n in l.loaded_networks), abs_sums=abs_sums) + return True + + def drop(layer_name): """Forget a layer's selection entry (its factors were removed or restored).""" if layer_name is not None and state['entries'].pop(layer_name, None) is not None: diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 366997118..f62fe431a 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -117,18 +117,20 @@ def network_activate(include=None, exclude=None): weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is not None and not isinstance(weights_backup, bool): network_apply_weights(module, None, None, device=device) - per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) - if sel_bias is None: - applied = lora_sdnq.apply_select(module, network_layer_name, per_net, component_wanted) - if applied is not None: - if applied and component_wanted: - applied_layers.append(network_layer_name) - applied_weight += 1 - module.network_current_names = component_wanted - module.network_current_stack = stack_sig - if task is not None: - pbar.update(task, advance=1) - continue + applied = lora_sdnq.apply_select_cached(module, network_layer_name, component_wanted) # a stored score record and factor pair serve before the deltas are assembled + if applied is None: + per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) + if sel_bias is None: + applied = lora_sdnq.apply_select(module, network_layer_name, per_net, component_wanted) + if applied is not None: + if applied and component_wanted: + applied_layers.append(network_layer_name) + applied_weight += 1 + module.network_current_names = component_wanted + module.network_current_stack = stack_sig + if task is not None: + pbar.update(task, advance=1) + continue lora_stack.warn_once('select-unridable', f'Network stack: mode={lora_stack.mode()} layer="{network_layer_name}" fallback=sum') # a pair the channel cannot carry (bias delta or malformed member) sums like any unsupported set elif getattr(module, 'sdnq_dequantizer', None) is not None: # hosting disabled: quantized layers have no side-channel to carry segments and packed backups cannot flip, so the sum paths below take the layer if any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks): @@ -137,8 +139,18 @@ def network_activate(include=None, exclude=None): sel_backup = network_backup_weights(module, network_layer_name, component_wanted, fuse) weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is not None and not isinstance(weights_backup, bool): + if lora_stack.register_weight_pair_cached(network_layer_name, module, component_wanted): # a stored score record registers without assembling the pair + backup_size += sel_backup + network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner + applied_layers.append(network_layer_name) + applied_weight += 1 + module.network_current_names = component_wanted + module.network_current_stack = stack_sig + if task is not None: + pbar.update(task, advance=1) + continue per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) - if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net): + if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net, component_wanted): backup_size += sel_backup # counted only when this branch keeps the layer; the fallthrough re-enters the shared backup call below, which counts it then network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner applied_layers.append(network_layer_name) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 7b4455077..61627f5f8 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -2003,6 +2003,122 @@ def test_flip_lands_before_crossover_step(): return True +def test_score_pair_chunked_precision(): + torch.manual_seed(71) + shapes = [(700, 460), (OUT_F, IN_F), (64,)] # off-chunk rows, square, and a 1-D norm delta + for shape in shapes: + d0 = (torch.randn(*shape, device=DEVICE) * 1e-2).to(torch.bfloat16) + d1 = (torch.randn(*shape, device=DEVICE) * 1e-2).to(torch.bfloat16) + for mode_name in ('klora', 'estlora'): + with select_mode(mode_name): + (s0, s1), (a0, a1) = lora_stack.score_pair(d0, d1, 8, 8) + f0, f1 = d0.to(torch.float64), d1.to(torch.float64) + ra0, ra1 = float(f0.abs().sum()), float(f1.abs().sum()) + if mode_name == 'klora': + k = 64 + rs0 = float(torch.topk(f0.abs().flatten(), min(k, f0.numel()), sorted=False).values.sum()) + rs1 = float(torch.topk(f1.abs().flatten(), min(k, f1.numel()), sorted=False).values.sum()) + else: + rs0, rs1 = float(f0.square().sum()), float(f1.square().sum()) + for got, ref, label in ((s0, rs0, 'score0'), (s1, rs1, 'score1'), (a0, ra0, 'abs0'), (a1, ra1, 'abs1')): + assert abs(got - ref) <= 1e-9 * max(abs(ref), 1e-12), f'{mode_name} {label} shape={shape}: got {got!r} ref {ref!r}' + frozen = torch.randn(300, 200, device=DEVICE, dtype=torch.float32) * 1e-2 # fp32 input aliases through to(); abs must stay out-of-place + pristine = frozen.clone() + with select_mode('klora'): + lora_stack.score_pair(frozen, frozen, 4, 4) + assert torch.equal(frozen, pristine), 'score_pair must not mutate a caller-owned fp32 delta' + return True + + +def test_select_replay_from_cache_skips_calc(): + import tempfile + layer = build_layer('uint4') + torch.manual_seed(73) + Dd1 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + Dd2 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + with tempfile.TemporaryDirectory() as tmp: + with host_rank(32), host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=layer), select_mode('klora'): + n1 = make_dense_net('selk1', layer, Dd1) + n2 = make_dense_net('selk2', layer, Dd2) + for net in (n1, n2): + lora_file = os.path.join(tmp, f'{net.name}.safetensors') + with open(lora_file, 'wb') as f: + f.write(b'0' * 64) + net.network_on_disk.filename = lora_file + from modules.modeldata import model_data + model_data.sd_model.sd_checkpoint_info = MockCheckpointInfo('test/cache-model') + Wdq0 = dq(layer) + with counting_calc() as calls: + activate(n1, n2) + assert calls['n'] > 0, 'a fresh select apply must assemble both deltas' + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None and entry['kind'] == 'factor' + fresh_scores, fresh_abs = entry['scores'], entry['abs_sums'] + first = dq(layer) + activate() + calls['n'] = 0 + real_svd = torch.svd_lowrank + torch.svd_lowrank = raise_no_svd + try: + activate(n1, n2) + finally: + torch.svd_lowrank = real_svd + assert calls['n'] == 0, f'a select replay must not assemble deltas: calc_updown ran {calls["n"]} times' + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None and entry['kind'] == 'factor', 'the replay must re-register the selection' + assert entry['scores'] == fresh_scores and entry['abs_sums'] == fresh_abs, 'cached scores must replay exactly' + assert torch.equal(dq(layer), first), 'select replay must be bit-identical to the fresh apply' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_select_weight_replay_from_cache_skips_calc(): + import tempfile + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device=DEVICE) + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.02) + lin.network_layer_name = 'lora_transformer_plainsel' + lin.network_current_names = () + torch.manual_seed(75) + Dd1 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + Dd2 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + W0 = lin.weight.detach().float().clone() + with tempfile.TemporaryDirectory() as tmp: + with host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=lin), select_mode('klora'): + n1 = make_dense_net('selw1', lin, Dd1) + n2 = make_dense_net('selw2', lin, Dd2) + for net in (n1, n2): + lora_file = os.path.join(tmp, f'{net.name}.safetensors') + with open(lora_file, 'wb') as f: + f.write(b'0' * 64) + net.network_on_disk.filename = lora_file + from modules.modeldata import model_data + model_data.sd_model.sd_checkpoint_info = MockCheckpointInfo('test/cache-model') + with counting_calc() as calls: + activate(n1, n2) + assert calls['n'] > 0, 'a fresh weight-kind select apply must assemble the pair' + entry = lora_stack.state['entries'].get('lora_transformer_plainsel') + assert entry is not None and entry['kind'] == 'weight' + fresh_scores = entry['scores'] + assert torch.equal(lin.weight.detach().float(), W0), 'weights stay pristine until the schedule applies a winner' + lora_stack.reset(20) + fresh_selected = lin.weight.detach().float().clone() + activate() + assert torch.equal(lin.weight.detach().float(), W0) + calls['n'] = 0 + activate(n1, n2) + assert calls['n'] == 0, f'a weight-kind select replay must not assemble the pair: calc_updown ran {calls["n"]} times' + entry = lora_stack.state['entries'].get('lora_transformer_plainsel') + assert entry is not None and entry['kind'] == 'weight', 'the replay must register from the score record' + assert entry['scores'] == fresh_scores, 'cached scores must replay exactly' + lora_stack.reset(20) # the winner recompute at schedule time still assembles its own delta, by design + assert torch.equal(lin.weight.detach().float(), fresh_selected), 'the replayed schedule must select the same winner' + activate() + assert torch.equal(lin.weight.detach().float(), W0) + return True + + CAT_COMPILE = category('compile') @@ -2194,7 +2310,8 @@ def run_tests(): test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled, test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer, - test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step]: + test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step, + test_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: From c762bfec18541c0a6d95f7eca2cff3c7caaab9d4 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 16:41:27 +0100 Subject: [PATCH 18/23] perf(lora): materialize select winners on the accelerator and time the reset The weight-kind schedule reset ran each winner's calc_updown on the target weight's device, which on a block-swapped denoiser is the cpu; at hundreds of layers per pass the cpu matmuls dominated every select generation. The delta now computes on the accelerator and moves back, matching the activate walk's convention. - reset and flip execution log a debug timing line (materialize, select loop, move/calc/apply split); the reset runs outside the activate walk, so its cost was invisible to the load timers --- modules/lora/lora_stack.py | 31 +++++++++++++++-- test/test-sdnq-lora-factors.py | 63 +++++++++++++++++++++++++++++++++- 2 files changed, 90 insertions(+), 4 deletions(-) diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index acd1117fa..85818c8b9 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -15,6 +15,7 @@ EST-LoRA arXiv:2508.02165 (its measured style-discrepancy estimate is exposed as an option instead of being derived from probe generations). """ +import time import weakref import hashlib @@ -305,9 +306,16 @@ def finalize(total_steps): e_den = sum(e['scores'][1] for e in state['entries'].values()) state['gamma_e'] = (e_num / e_den) if e_den > 0 else 1.0 state['flips'] = {} - if any(e['kind'] == 'weight' for e in state['entries'].values()): + stats = {'weight_n': 0, 'factor_n': 0, 'materialize': 0.0, 'select': 0.0, 'w_move': 0.0, 'w_calc': 0.0, 'w_apply': 0.0} + state['stats'] = stats + stats['weight_n'] = sum(1 for e in state['entries'].values() if e['kind'] == 'weight') + stats['factor_n'] = len(state['entries']) - stats['weight_n'] + if stats['weight_n'] > 0: + t0 = time.time() materialize_model() + stats['materialize'] = time.time() - t0 style_first = 0 + t0 = time.time() for layer_name, entry in list(state['entries'].items()): # snapshot: apply_selection drops entries whose module died flip_at = layer_flip_step(entry['scores'], state['total_steps']) initial = 1 if flip_at == 0 else 0 @@ -315,6 +323,7 @@ def finalize(total_steps): apply_selection(layer_name, entry, initial) if 0 < flip_at < state['total_steps']: state['flips'].setdefault(flip_at - 1, []).append(layer_name) # step callbacks fire after the denoise, so the flip runs one step early to be live during the crossover step's forward + stats['select'] = time.time() - t0 state['finalized'] = True if len(state['entries']) > 0: # only a built schedule can carry a flip count, so this is the line that shows selection is live rather than requested gamma = state['gamma_e'] if mode() == 'estlora' else state['gamma'] @@ -322,6 +331,8 @@ def finalize(total_steps): if report != state['reported']: # rebuilt every pass, so a batch would otherwise repeat one line per image state['reported'] = report log.info(f'Network load: type=LoRA stack={report[0]} layers={report[1]} style={report[2]} flips={report[3]} steps={report[4]} gamma={report[5]:.3f}') + # logged every pass: the reset runs outside the activate walk, so its cost is invisible to the load timers + log.debug(f'Network select: type=LoRA reset weight={stats["weight_n"]} factor={stats["factor_n"]} time={{materialize: {stats["materialize"]:.2f}, select: {stats["select"]:.2f}, move: {stats["w_move"]:.2f}, calc: {stats["w_calc"]:.2f}, apply: {stats["w_apply"]:.2f}}}') def reset(total_steps): @@ -335,10 +346,15 @@ def on_step(step): """Flip the layers whose crossover is this step; non-flip steps are a dict miss.""" if not state['finalized']: return - for layer_name in state['flips'].get(int(step), ()): + layers = state['flips'].get(int(step), ()) + if not layers: + return + t0 = time.time() + for layer_name in layers: entry = state['entries'].get(layer_name) if entry is not None: apply_selection(layer_name, entry, 1) + log.debug(f'Network select: type=LoRA flip step={int(step)} layers={len(layers)} time={time.time() - t0:.2f}') def apply_selection(layer_name, entry, winner): @@ -375,6 +391,15 @@ def weight_selection(module, entry, winner): if weight is None or weight.is_meta: warn_once('select-offloaded', 'Network stack: flip=skipped weight=offloaded') return + from modules import devices + stats = state.get('stats') or {} device = weight.device - updown = net_module.calc_updown(backup.to(device))[0] + t0 = time.time() + base = backup.to(devices.device) # a swapped-out layer keeps its weight on cpu; the delta matmul belongs on the accelerator regardless + t1 = time.time() + updown = net_module.calc_updown(base)[0].to(device) + t2 = time.time() network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it + stats['w_move'] = stats.get('w_move', 0.0) + (t1 - t0) + stats['w_calc'] = stats.get('w_calc', 0.0) + (t2 - t1) + stats['w_apply'] = stats.get('w_apply', 0.0) + (time.time() - t2) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 61627f5f8..e5c4f82fb 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -2119,6 +2119,66 @@ def test_select_weight_replay_from_cache_skips_calc(): return True +def test_select_reset_reports_timing(): + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device=DEVICE) + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.02) + lin.network_layer_name = 'lora_transformer_timed' + lin.network_current_names = () + A1, B1, _D1 = make_delta(seed=77, sigma=1e-2) + A2, B2, _D2 = make_delta(seed=78, sigma=1e-2) + n1 = make_net('t1', lin, A1, B1) + n2 = make_net('t2', lin, A2, B2) + with mock_model(lin=lin), select_mode('klora'): + activate(n1, n2) + lora_stack.reset(20) + stats = lora_stack.state.get('stats') + assert stats is not None, 'a reset must publish its timing stats' + assert stats['weight_n'] == 1 and stats['factor_n'] == 0, f'weight-kind counts wrong: {stats}' + assert stats['w_calc'] > 0.0, 'the weight-kind winner apply must account its calc time' + assert stats['select'] > 0.0 + activate() + layer = build_layer('uint4') + f1, f2, _Df1, _Df2 = select_pair(layer, seed0=79, seed1=80) + with mock_model(lin=layer), select_mode('klora'): + activate(f1, f2) + lora_stack.reset(20) + stats = lora_stack.state.get('stats') + assert stats is not None and stats['factor_n'] == 1 and stats['weight_n'] == 0, f'factor-kind counts wrong: {stats}' + assert stats['w_calc'] == 0.0, 'factor-kind resets flip segments and must not touch the weight path' + activate() + return True + + +def test_select_weight_flip_calcs_on_accelerator(): + from modules.lora import network_lora + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device='cpu') # a swapped-out layer: weight lives on cpu + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F) * 0.02) + lin.network_layer_name = 'lora_transformer_swapped' + lin.network_current_names = () + A1, B1, _D1 = make_delta(seed=81, sigma=1e-2) + A2, B2, _D2 = make_delta(seed=82, sigma=1e-2) + n1 = make_net('s1', lin, A1, B1) + n2 = make_net('s2', lin, A2, B2) + seen = [] + real = network_lora.NetworkModuleLora.calc_updown + def spy(self, target, *args, **kwargs): + seen.append(target.device.type) + return real(self, target, *args, **kwargs) + with mock_model(lin=lin), select_mode('klora'): + activate(n1, n2) + lin.to('cpu') # the offload dispatch swaps blocks back out after the walk; the reset must not follow the weight onto the cpu + network_lora.NetworkModuleLora.calc_updown = spy + try: + lora_stack.reset(20) + finally: + network_lora.NetworkModuleLora.calc_updown = real + assert seen and all(d == DEVICE.type for d in seen), f'winner materialization must calc on the accelerator, saw {seen}' + activate() + return True + + CAT_COMPILE = category('compile') @@ -2311,7 +2371,8 @@ def run_tests(): test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer, test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step, - test_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc]: + test_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc, + test_select_reset_reports_timing, test_select_weight_flip_calcs_on_accelerator]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: From 455a6d0e8f41a94cbd299293fb76aeae8fee7b1c Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 20:20:43 +0100 Subject: [PATCH 19/23] feat(lora): default the select-mode ramp to 0 At 1.5 the ramp hands nearly every layer to the style network mid-generation, which on few-step models overwrites the forming subject before identity sets. Alpha 0 freezes the schedule into a static per-layer split: the subject keeps its layers for the whole generation and style keeps the layers where it is more salient. Nonzero values remain the scheduled handover toward the second network. - locale hints describe both regimes; the mode hint no longer implies the shift is always on --- modules/ui_definitions.py | 2 +- ui/locale/locale_en.json | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 1808f62a7..d52a142b6 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -711,7 +711,7 @@ def create_settings(cmd_opts): "lora_stack_sep": OptionInfo("

Stacking options

", "", gr.HTML), "lora_stack_mode": OptionInfo("sum", "LoRA stack mode", gr.Dropdown, {"choices": ["sum", "ties", "dare_ties", "dare_linear", "magnitude_prune", "klora", "estlora"]}), "lora_stack_density": OptionInfo(0.5, "LoRA stack density", gr.Slider, {"minimum": 0.05, "maximum": 1.0, "step": 0.05}), - "lora_stack_alpha": OptionInfo(1.5, "LoRA stack ramp", gr.Slider, {"minimum": 0.0, "maximum": 3.0, "step": 0.1}), + "lora_stack_alpha": OptionInfo(0.0, "LoRA stack ramp", gr.Slider, {"minimum": 0.0, "maximum": 3.0, "step": 0.1}), "lora_stack_discrepancy": OptionInfo(0.5, "LoRA stack discrepancy", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}), "lora_meta_sep": OptionInfo("

Metadata

", "", gr.HTML), diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index f52b58a91..3fb030aa1 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -867,9 +867,9 @@ {"id":"","label":"LoRA quantized host rank","localized":"","hint":"Maximum rank used to carry adapter types that are not natively low-rank (LoKR, LoHA, OFT, DoRA) alongside the quantized weights instead of merging them in.
Higher values retain more of the adapter at proportionally more memory. Plain LoRA files are carried exactly at their own rank.

Applies only to SDNQ models quantized below 8 bits, where merging erases most of the adapter; at 8 bits and above merging retains it and hosting is skipped.

0 disables hosting and merges every adapter into the quantized weights.

Default is 256.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized host calibration","localized":"","hint":"Collects per-channel activation statistics from the model's own generations and uses them to focus hosted-adapter truncation on the channels with the strongest activations.
Statistics accumulate in the background on models quantized below 8 bits, persist per checkpoint, and raise delivered adapter fidelity at the same LoRA quantized host rank, most at low ranks.

Capture is skipped while the model is compiled; previously cached statistics still apply.

Enabled by default.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized host cache","localized":"","hint":"Disk space in GB for caching computed hosting factors.
A cached set skips the truncation math on the next load; least recently used entries are evicted once the budget is exceeded.

0 disables the cache.

Default is 10.","ui":"settings_lora"}, - {"id":"","label":"LoRA stack mode","localized":"","hint":"How multiple networks targeting the same layer are combined:
- sum: adds all contributions
- ties: keeps each network's strongest elements and merges only where signs agree
- dare_ties: randomly drops elements, rescales the survivors, then merges where signs agree
- dare_linear: randomly drops elements, rescales the survivors and sums
- magnitude_prune: keeps each network's strongest elements and sums
- klora / estlora: assign each layer to one of exactly two networks, the first in the prompt as subject and the second as style, shifting from subject toward style over the sampling steps

Each layer is given to a single network at a time, so a subject and a style that both need sustained strength can end up under-applied. For reliable blending of two strong networks, sum, ties and dare_ties apply every network throughout and combine more fully.

Kept fractions are set by LoRA stack density; the subject-to-style shift by LoRA stack ramp and LoRA stack discrepancy.

Applies to the native load path; other load methods and text encoder networks always combine as sum. Selection modes fall back to sum unless exactly two networks are loaded, or when model compile is active.

Default is sum.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack mode","localized":"","hint":"How multiple networks targeting the same layer are combined:
- sum: adds all contributions
- ties: keeps each network's strongest elements and merges only where signs agree
- dare_ties: randomly drops elements, rescales the survivors, then merges where signs agree
- dare_linear: randomly drops elements, rescales the survivors and sums
- magnitude_prune: keeps each network's strongest elements and sums
- klora / estlora: assign each layer to one of exactly two networks, the first in the prompt as subject and the second as style; LoRA stack ramp optionally shifts layers toward style over the sampling steps

Each layer is given to a single network at a time, so a subject and a style that both need sustained strength can end up under-applied. For reliable blending of two strong networks, sum, ties and dare_ties apply every network throughout and combine more fully.

Kept fractions are set by LoRA stack density; the subject-to-style shift by LoRA stack ramp and LoRA stack discrepancy.

Applies to the native load path; other load methods and text encoder networks always combine as sum. Selection modes fall back to sum unless exactly two networks are loaded, or when model compile is active.

Default is sum.","ui":"settings_lora"}, {"id":"","label":"LoRA stack density","localized":"","hint":"Fraction of elements each network keeps under the ties, dare_ties, dare_linear and magnitude_prune stack modes.
Lower values keep only the strongest contributions and reduce interference between networks at the cost of per-network detail. The dare variants drop at random and rescale the survivors to preserve expected strength.

Default is 0.5.","ui":"settings_lora"}, - {"id":"","label":"LoRA stack ramp","localized":"","hint":"Slope of the subject-to-style shift across the sampling steps in the klora and estlora stack modes.
Higher values shift layers to the style network earlier and more broadly; lower values keep the subject network dominant for longer.

0 keeps the balance fixed for the whole generation.

Default is 1.5.","ui":"settings_lora"}, + {"id":"","label":"LoRA stack ramp","localized":"","hint":"Slope of the subject-to-style shift across the sampling steps in the klora and estlora stack modes.

0 keeps the layer assignment fixed for the whole generation: each layer stays with the network that is more salient there, which preserves the subject while the style keeps its own layers. Higher values hand layers to the style network progressively, ending in a style takeover; on few-step models the handover happens early enough to override the subject.

Default is 0.","ui":"settings_lora"}, {"id":"","label":"LoRA stack discrepancy","localized":"","hint":"Stand-in for the measured style separation the estlora stack mode would otherwise derive from data.
Higher values keep layers with the subject network longer; lower values let the style network take layers earlier.

Layer scores are balanced by each network's overall strength, so a louder network does not take layers on magnitude alone.

Applies only when LoRA stack mode is estlora.

Default is 0.5.","ui":"settings_lora"}, {"id":"","label":"LoRA auto-apply tags","localized":"","hint":"Automatically add trigger words/tags from LoRA metadata to your prompt.
Set to the number of tags to auto-apply, e.g., 3 = add top 3 trigger tags.
Set to 0 to disable, -1 to add all available tags.","ui":"settings_lora"}, {"id":"","label":"LoRA memory cache","localized":"","hint":"How many LoRAs to keep in network for future use before requiring reloading from storage","ui":"settings_lora"}, From 2620b0cc5b83c15070e20f3f038a5e9008dfe4fe Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 21:34:28 +0100 Subject: [PATCH 20/23] feat(lora): per-block strength scales each targeted layer's delta by a slot of a per-architecture block vector. VALUE is a preset name, a scalar, or a comma vector; presets stretch onto the block count of the current model and the a1111 17-slot and 12-slot layouts are accepted on sd and sdxl. The factor enters through the module multiplier, so every apply path carries it: the exact factor channel, hosting, requantize routing, dense stack combines and select scoring. - modules/lora/lora_blocks.py: slot classification from network_layer_mapping (namespace-first, anchored chain prefixes), preset resolution reusing the merge block-weight tables with BASE forced neutral, generated classic segment names plus DOUBLE/SINGLE chain names, per-model memoization - the raw spec stages through pending_config and promotes with the other multipliers, keeping fuse removal consistent - block weights join the activation signature, the per-module apply stamp and the factor cache identity; entries without block weights keep their existing signature bytes - non-native load methods warn once and ignore the argument --- modules/lora/extra_networks_lora.py | 17 +- modules/lora/lora_blocks.py | 348 ++++++++++++++++++++++++ modules/lora/lora_factor_cache.py | 7 +- modules/lora/lora_load.py | 3 +- modules/lora/network.py | 21 +- modules/lora/networks.py | 4 +- test/test-sdnq-lora-factors.py | 396 +++++++++++++++++++++++++++- 7 files changed, 778 insertions(+), 18 deletions(-) create mode 100644 modules/lora/lora_blocks.py diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index c0fce6fb4..81015af01 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -98,6 +98,7 @@ def parse(p, params_list, step=0): unet_multipliers = [] dyn_dims = [] lora_modules = [] + block_specs = [] for params in params_list: name = params.positional[0] @@ -131,6 +132,7 @@ def parse(p, params_list, step=0): te_multipliers.append(te_multiplier) unet_multipliers.append(unet_multiplier) dyn_dims.append(dyn_dim) + block_specs.append(params.named.get('lbw', None)) # per-block strength; resolved per layer by lora_blocks lora_module = [] name_lower = params.positional[0].lower() @@ -150,7 +152,7 @@ def parse(p, params_list, step=0): lora_modules.append(lora_module) - return names, te_multipliers, unet_multipliers, dyn_dims, lora_modules + return names, te_multipliers, unet_multipliers, dyn_dims, lora_modules, block_specs def unload_diffusers(): @@ -174,8 +176,9 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): self.model = None self.errors = {} - def signature(self, names: list[str], te_multipliers: list, unet_multipliers: list): - return [f'{name}:{te}:{unet}' for name, te, unet in zip(names, te_multipliers, unet_multipliers, strict=False)] + def signature(self, names: list[str], te_multipliers: list, unet_multipliers: list, block_specs: list | None = None): + specs = block_specs if block_specs else [None] * len(names) + return [f'{name}:{te}:{unet}' + (f':lbw={str(spec).strip().lower()}' if spec else '') for name, te, unet, spec in zip(names, te_multipliers, unet_multipliers, specs, strict=False)] def changed(self, requested: list[str], include: list[str] | None = None, exclude: list[str] | None = None) -> tuple[bool, str]: from modules.lora import lora_sdnq, lora_stack @@ -220,14 +223,16 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): if len(params_list) > 0 and not self.active: # activate patches once self.active = True self.model = shared.opts.sd_model_checkpoint - names, te_multipliers, unet_multipliers, dyn_dims, lora_modules = parse(p, params_list, step) - requested = self.signature(names, te_multipliers, unet_multipliers) + names, te_multipliers, unet_multipliers, dyn_dims, lora_modules, block_specs = parse(p, params_list, step) + requested = self.signature(names, te_multipliers, unet_multipliers, block_specs) reason = '' load_method, load_reason = lora_overrides.get_method() from modules.lora import lora_stack if load_method != 'native' and lora_stack.mode() != 'sum': log.warning(f'Network stack: mode={lora_stack.mode()} method={load_method} fallback=sum') + if load_method != 'native' and any(block_specs): + log.warning(f'Network blocks: method={load_method} fallback=none') if debug: import sys fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access @@ -254,7 +259,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): has_changed = lora_nunchaku.load_nunchaku(names, unet_multipliers) else: # native - lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims, activate=False) # load only, activation below honors include/exclude + lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims, block_specs=block_specs, activate=False) # load only, activation below honors include/exclude has_changed, reason = self.changed(requested, include, exclude) if has_changed: jobid = shared.state.begin('LoRA') diff --git a/modules/lora/lora_blocks.py b/modules/lora/lora_blocks.py new file mode 100644 index 000000000..211f1ac78 --- /dev/null +++ b/modules/lora/lora_blocks.py @@ -0,0 +1,348 @@ +"""Per-block LoRA strength: . + +Each targeted layer maps to one slot of a per-architecture weight vector and +the network's multiplier is scaled by that slot. Slot 0 is BASE: on unet +architectures it covers the text encoder and the unet layers outside the +block chain, on transformer architectures the layers outside the block +chain(s). The remaining slots follow the merge block-weight layout on unet +architectures (26 on sd, 20 on sdxl: input blocks, mid, output blocks) and +the transformer chain(s) in depth order elsewhere, with chain lengths +scanned from the live network_layer_mapping rather than hardcoded. + +VALUE is a preset name (case-insensitive), a single number broadcast to +every slot, or a comma list with one number per slot. Named presets force +BASE to 1.0, since the merge tables carry 0 there with merge semantics, and +stretch onto the block count of the current model; classic segment names +(INS, OUTALL, ...) generate from ranges, so they also work on transformer +chains via thirds, and DOUBLE/SINGLE mute one chain on two-chain +architectures. Explicit vectors are taken verbatim at the slot count, with +the a1111 17-slot (sd) and 12-slot (sdxl) layouts accepted and expanded, +omitted slots neutral. A value that fits nothing is ignored with a warning +and the network applies at its plain strength. +""" + +import re + +from modules import shared +from modules.logger import log +from modules.lora import lora_common as l + + +UNET_ARCHES = ('sd', 'sdxl') +CHAINS = { # arch -> anchored tail prefixes, one per chain, in depth order + 'sd3': ('transformer_blocks_',), + 'anima': ('transformer_blocks_',), + 'f1': ('transformer_blocks_', 'single_transformer_blocks_'), + 'f2': ('transformer_blocks_', 'single_transformer_blocks_'), + 'chroma': ('transformer_blocks_', 'single_transformer_blocks_'), + 'zimage': ('layers_',), + 'ernieimage': ('layers_',), + 'krea2': ('blocks_',), +} +CLASSIC = ('ALL', 'NONE', 'INALL', 'INS', 'IND', 'MIDD', 'OUTALL', 'OUTD', 'OUTS') +CHAIN_NAMES = ('DOUBLE', 'SINGLE') +SD1_17 = (0, 2, 3, 5, 6, 8, 9, 13, 17, 18, 19, 20, 21, 22, 23, 24, 25) # BASE, IN01, IN02, IN04, IN05, IN07, IN08, MID, OUT03..OUT11 +SDXL_12 = (0, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15, 16) # BASE, IN04, IN05, IN07, IN08, MID, OUT00..OUT05 +VECTOR_MEMO_CAP = 64 +MISS = object() + +re_down = re.compile(r'^down_blocks_(\d+)_(resnets|attentions|downsamplers)_(\d+)') +re_up = re.compile(r'^up_blocks_(\d+)_(resnets|attentions|upsamplers)_(\d+)') +re_chain_index = re.compile(r'^(\d+)') + +state: dict = {'stamp': None, 'layout': None, 'index': {}, 'vectors': {}} +warned: set = set() + + +def warn_once(key, message): + if key not in warned: + warned.add(key) + log.warning(message) + + +def build_unet_layout(arch, mapping): + down, up = -1, -1 + for key in mapping: + if not key.startswith('lora_unet_'): + continue + tail = key[len('lora_unet_'):] + m = re_down.match(tail) + if m is not None: + down = max(down, int(m.group(1))) + continue + m = re_up.match(tail) + if m is not None: + up = max(up, int(m.group(1))) + if down < 0 or up < 0: + return None + n_in = 3 * (down + 1) # conv_in plus two pairs and a sampler slot per group: the compvis input_blocks count + n_out = 3 * (up + 1) + n = 2 + n_in + n_out + return { + 'arch': arch, 'kind': 'unet', 'n': n, 'n_in': n_in, + 'ins': list(range(1, 1 + n_in)), + 'mids': [1 + n_in], + 'outs': list(range(2 + n_in, n)), + } + + +def build_dit_layout(arch, mapping): + prefixes = CHAINS.get(arch) + if prefixes is None: + return None + counts = [0 for _ in prefixes] + for key in mapping: + if not key.startswith('lora_transformer_'): + continue + tail = key[len('lora_transformer_'):] + for i, prefix in enumerate(prefixes): + if tail.startswith(prefix): + m = re_chain_index.match(tail[len(prefix):]) + if m is not None: + counts[i] = max(counts[i], int(m.group(1)) + 1) + break + total = sum(counts) + if total == 0: + return None + chains = [] + offset = 0 + for prefix, count in zip(prefixes, counts, strict=False): + chains.append((prefix, count, offset)) + offset += count + n = 1 + total + blocks = list(range(1, n)) + return { + 'arch': arch, 'kind': 'dit', 'n': n, 'chains': chains, + 'ins': [s for i, s in enumerate(blocks) if i * 3 // total == 0], + 'mids': [s for i, s in enumerate(blocks) if i * 3 // total == 1], + 'outs': [s for i, s in enumerate(blocks) if i * 3 // total == 2], + } + + +def layout(): + sd_model = getattr(shared, 'sd_model', None) + mapping = getattr(sd_model, 'network_layer_mapping', None) if sd_model is not None else None + if not mapping: + return None + arch = shared.sd_model_type + stamp = (arch, id(mapping)) + if state['stamp'] == stamp: + return state['layout'] + state['stamp'] = stamp + state['layout'] = build_unet_layout(arch, mapping) if arch in UNET_ARCHES else build_dit_layout(arch, mapping) + state['index'].clear() + state['vectors'].clear() + return state['layout'] + + +def classify(sd_key, lay): + if sd_key.startswith('lora_te'): + return 0 if lay['kind'] == 'unet' else None # BASE covers the TE on unet arches; transformer vectors do not model the TE + if sd_key.startswith('lora_llm_adapter_'): + return None + if lay['kind'] == 'unet': + if not sd_key.startswith('lora_unet_'): + return None + tail = sd_key[len('lora_unet_'):] + m = re_down.match(tail) + if m is not None: + slot = 1 + 3 * int(m.group(1)) + (2 if m.group(2) == 'downsamplers' else int(m.group(3))) + return 1 + slot + m = re_up.match(tail) + if m is not None: + slot = 3 * int(m.group(1)) + (2 if m.group(2) == 'upsamplers' else int(m.group(3))) + return 2 + lay['n_in'] + slot + if tail.startswith('mid_block'): + return 1 + lay['n_in'] + if tail.startswith('conv_in'): + return 1 # IN00 + if tail.startswith('conv_out') or tail.startswith('conv_norm_out'): + return lay['n'] - 1 # the compvis out group belongs to the last output block + return 0 # time_embedding, add_embedding and other non-block leaves + if not sd_key.startswith('lora_transformer_'): + return None + tail = sd_key[len('lora_transformer_'):] + for prefix, count, offset in lay['chains']: + if tail.startswith(prefix): + m = re_chain_index.match(tail[len(prefix):]) + if m is not None and int(m.group(1)) < count: + return 1 + offset + int(m.group(1)) + return 0 + return 0 # embedders, projections, refiners and other non-chain layers + + +def block_index(sd_key): + lay = layout() + if lay is None: + return None + cached = state['index'].get(sd_key, MISS) + if cached is not MISS: + return cached + idx = classify(sd_key, lay) + state['index'][sd_key] = idx + return idx + + +def fill_band(vec, slots, lo, hi): + k = len(slots) + for i, s in enumerate(slots): + if lo * k <= i < hi * k: + vec[s] = 1.0 + + +def classic_vector(name, lay): + if name == 'ALL': + return [1.0] * lay['n'] + vec = [0.0] * lay['n'] + if name == 'NONE': + return vec + vec[0] = 1.0 + if name == 'INALL': + fill_band(vec, lay['ins'], 0.0, 1.0) + elif name == 'INS': # shallow half of the input side + fill_band(vec, lay['ins'], 0.0, 0.5) + elif name == 'IND': # deep half of the input side + fill_band(vec, lay['ins'], 0.5, 1.0) + elif name == 'MIDD': # the middle of the network: deep input half, mid, deep output half + fill_band(vec, lay['ins'], 0.5, 1.0) + fill_band(vec, lay['mids'], 0.0, 1.0) + fill_band(vec, lay['outs'], 0.0, 0.5) + elif name == 'OUTALL': + fill_band(vec, lay['outs'], 0.0, 1.0) + elif name == 'OUTD': # deep half of the output side, nearest the mid + fill_band(vec, lay['outs'], 0.0, 0.5) + elif name == 'OUTS': # shallow half of the output side, nearest the image + fill_band(vec, lay['outs'], 0.5, 1.0) + return vec + + +def chain_vector(name, lay): + chains = lay.get('chains') or [] + if len(chains) != 2: + return None + vec = [1.0] * lay['n'] + keep = 0 if name == 'DOUBLE' else 1 + for i, (_prefix, count, offset) in enumerate(chains): + val = 1.0 if i == keep else 0.0 + for s in range(1 + offset, 1 + offset + count): + vec[s] = val + return vec + + +def stretch(src, k): + if k == len(src): + return [float(v) for v in src] + out = [] + for i in range(k): + x = i * (len(src) - 1) / (k - 1) if k > 1 else 0.0 + lo = int(x) + hi = min(lo + 1, len(src) - 1) + f = x - lo + out.append(float(src[lo]) * (1.0 - f) + float(src[hi]) * f) + return out + + +def preset_vector(name, lay): + from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS + if name in CHAIN_NAMES: + return chain_vector(name, lay) + if name in CLASSIC: + return classic_vector(name, lay) + if lay['arch'] == 'sdxl': + src = SDXL_BLOCK_WEIGHTS_PRESETS.get(name) or SDXL_BLOCK_WEIGHTS_PRESETS.get('SDXL_' + name) + if src is not None: + return [1.0] + [float(v) for v in src[1:]] # merge tables carry 0 in the BASE slot; a preset must leave the TE alone + if name.startswith('SDXL_'): + return None # explicitly arch-tagged, not reinterpreted elsewhere + src = BLOCK_WEIGHTS_PRESETS.get(name) + if src is None: + return None + if lay['arch'] == 'sd': + return [1.0] + [float(v) for v in src[1:]] + return [1.0] + stretch(src[1:], lay['n'] - 1) + + +def parse_vector(parts, lay): + try: + vals = [float(x) for x in parts] + except ValueError: + return None + n = lay['n'] + if len(vals) == n: + return vals + if len(vals) == n - 1: + return [1.0] + vals + legacy = SD1_17 if lay['arch'] == 'sd' else (SDXL_12 if lay['arch'] == 'sdxl' else None) + if legacy is not None and len(vals) == len(legacy): + vec = [1.0] * n # slots the a1111 layouts omit stay neutral + for slot, v in zip(legacy, vals, strict=False): + vec[slot] = v + return vec + return None + + +def resolve(spec): + """Resolve a raw lbw value into a slot vector for the current model, or None when it fits nothing.""" + lay = layout() + if lay is None: + return None + raw = str(spec).strip() + key = raw.lower() + if key in state['vectors']: + return state['vectors'][key] + if len(state['vectors']) > VECTOR_MEMO_CAP: + state['vectors'].clear() + vec = None + if ',' in raw: + vec = parse_vector([x.strip() for x in raw.split(',')], lay) + if vec is None: + warn_once(f'lbw-vector:{key}:{lay["arch"]}', f'Network blocks: value="{raw}" arch={lay["arch"]} expected={lay["n"]} fallback=none') + else: + try: + vec = [float(raw)] * lay['n'] + except ValueError: + vec = preset_vector(raw.upper(), lay) + if vec is None: + warn_once(f'lbw-name:{key}:{lay["arch"]}', f'Network blocks: preset="{raw}" arch={lay["arch"]} fallback=none') + if vec is not None: + log.info(f'Network blocks: value="{raw}" arch={lay["arch"]} slots={lay["n"]} range={min(vec):.2f}-{max(vec):.2f}') + state['vectors'][key] = vec + return vec + + +def factor(sd_key, net): + """Per-layer scale from a network's block vector; 1.0 whenever the vector does not apply.""" + try: + spec = getattr(net, 'block_spec', None) + if not spec: + return 1.0 + vec = resolve(spec) + if vec is None: + return 1.0 + idx = block_index(sd_key) + if idx is None: + return 1.0 + return float(vec[idx]) + except Exception as e: + warn_once('lbw-error', f'Network blocks: {e} fallback=none') + return 1.0 + + +def net_signature(net): + """Normalized spec of one network, or None; joins content identities such as the factor cache signature.""" + spec = getattr(net, 'block_spec', None) + if not spec: + return None + return str(spec).strip().lower() + + +def active(): + return any(getattr(net, 'block_spec', None) for net in l.loaded_networks) + + +def signature(): + """Identity suffix for the per-module apply stamp; empty while no loaded network carries block weights.""" + specs = [f'{net.name}:{net_signature(net)}' for net in l.loaded_networks if getattr(net, 'block_spec', None)] + if len(specs) == 0: + return '' + return '|lbw=' + ','.join(specs) diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index aab36c7a5..1da7106c6 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -59,6 +59,7 @@ def signature(wanted_names): 'stack': lora_stack.signature(), 'nets': [], } + from modules.lora import lora_blocks for name, te, unet, dyn in wanted_names: net = next((n for n in l.loaded_networks if n.name == name), None) filename = getattr(getattr(net, 'network_on_disk', None), 'filename', None) @@ -66,7 +67,11 @@ def signature(wanted_names): st = os.stat(filename) except Exception: return None - parts['nets'].append([name, repr(te), repr(unet), repr(dyn), filename, int(st.st_mtime), st.st_size]) + entry = [name, repr(te), repr(unet), repr(dyn), filename, int(st.st_mtime), st.st_size] + spec = lora_blocks.net_signature(net) + if spec is not None: # appended only when set so existing cache files stay valid without block weights + entry.append(spec) + parts['nets'].append(entry) return parts diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index b297108bd..f9a885cf2 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -261,7 +261,7 @@ def gather_networks(names): return networks_on_disk -def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None, lora_modules=None, activate=True): +def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None, lora_modules=None, block_specs=None, activate=True): networks_on_disk = gather_networks(names) failed_to_load_networks = [] recompile_model, skip_lora_load = maybe_recompile_model(names, te_multipliers) @@ -309,6 +309,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non 'te': te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier, 'unet': unet_multipliers[i] if unet_multipliers else shared.opts.extra_networks_default_multiplier, 'dyn': dyn_dims[i] if dyn_dims else None, # a multiplier is not a rank; float dyn_dim crashes every consumer that slices with it + 'blocks': block_specs[i] if block_specs and len(block_specs) > i else None, } l.loaded_networks.append(net) diff --git a/modules/lora/network.py b/modules/lora/network.py index b16060563..fed18dd2e 100644 --- a/modules/lora/network.py +++ b/modules/lora/network.py @@ -148,6 +148,7 @@ class Network: # LoraModule self.te_multiplier = 1.0 self.unet_multiplier = [1.0] * 3 self.dyn_dim = None + self.block_spec = None # raw lbw= value; per-layer factors resolve through lora_blocks self.pending_config = None # staged multipliers; network_activate promotes them after the removal pass so fuse removal subtracts the delta that was applied self.modules = {} self.mismatch = 0 # deltas dropped for not fitting their target module; try_load_chain refuses the file when non-zero @@ -195,15 +196,19 @@ class NetworkModule: def multiplier(self): unet_multiplier = 3 * [self.network.unet_multiplier] if not isinstance(self.network.unet_multiplier, list) else self.network.unet_multiplier if self.sd_key.startswith('lora_te') or 'transformer' in self.sd_key[:20]: - return self.network.te_multiplier - if "down_blocks" in self.sd_key: - return unet_multiplier[0] - if "mid_block" in self.sd_key: - return unet_multiplier[1] - if "up_blocks" in self.sd_key: - return unet_multiplier[2] + base = self.network.te_multiplier + elif "down_blocks" in self.sd_key: + base = unet_multiplier[0] + elif "mid_block" in self.sd_key: + base = unet_multiplier[1] + elif "up_blocks" in self.sd_key: + base = unet_multiplier[2] else: - return unet_multiplier[0] + base = unet_multiplier[0] + if getattr(self.network, 'block_spec', None) is None: # per-block strength is off for this network; no shared access on this path + return base + from modules.lora import lora_blocks + return base * lora_blocks.factor(self.sd_key, self.network) def calc_scale(self): if self.scale is not None: diff --git a/modules/lora/networks.py b/modules/lora/networks.py index f62fe431a..479ea998a 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -2,6 +2,7 @@ from contextlib import nullcontext import time import rich.progress as rp from modules.errorlimiter import limit_errors +from modules.lora import lora_blocks from modules.lora import lora_common as l from modules.lora import lora_overrides from modules.lora import lora_sdnq @@ -55,6 +56,7 @@ def network_activate(include=None, exclude=None): net.te_multiplier = pending['te'] net.unet_multiplier = pending['unet'] net.dyn_dim = pending['dyn'] + net.block_spec = pending.get('blocks', None) t0 = time.time() fuse = lora_overrides.fuse_native() # resolve once: backup and apply passes must agree with limit_errors("network_activate") as elimit: @@ -90,7 +92,7 @@ def network_activate(include=None, exclude=None): refused = 0 with devices.inference_context(), pbar: wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () - stack_sig = lora_stack.signature() + lora_sdnq.signature() # tracked beside network_current_names so settings-only stack or mechanism changes re-apply + stack_sig = lora_stack.signature() + lora_blocks.signature() + lora_sdnq.signature() # tracked beside network_current_names so settings-only stack, block-weight or mechanism changes re-apply select_active = len(l.loaded_networks) > 0 and lora_stack.active_select(len(l.loaded_networks)) # restore-only walks have nothing to stack; the count warning would fire on every network-free generation applied_layers.clear() lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index e5c4f82fb..c0cf7fede 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -76,7 +76,7 @@ modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) from modules.errors import log # pylint: disable=wrong-import-position from modules import shared, sd_models # pylint: disable=wrong-import-position -from modules.lora import network, network_lora, lora_sdnq, lora_stack, networks # pylint: disable=wrong-import-position +from modules.lora import network, network_lora, lora_blocks, lora_sdnq, lora_stack, networks # pylint: disable=wrong-import-position from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position from sdnq.quantizer import sdnq_quantize_layer, SDNQConfig # pylint: disable=wrong-import-position @@ -2321,6 +2321,392 @@ def test_stacked_shape_mismatch_falls_back(): return True +# ============================================================ +# Tests - per-block strength (lbw) +# ============================================================ + +CAT_BLOCKS = category('block-weights') + + +def block_fixture_keys(arch): + """Sparse network_layer_mapping keys per arch: the layout scan only needs each chain's max index.""" + if arch == 'sd': + return ['lora_unet_down_blocks_3_resnets_1_conv1', 'lora_unet_up_blocks_3_resnets_2_conv1'] + if arch == 'sdxl': + return ['lora_unet_down_blocks_2_resnets_1_conv1', 'lora_unet_up_blocks_2_resnets_2_conv1'] + if arch in ('f1', 'chroma'): + return ['lora_transformer_transformer_blocks_18_attn_to_q', 'lora_transformer_single_transformer_blocks_37_attn_to_q'] + if arch == 'krea2': + return ['lora_transformer_blocks_27_attn_wq', 'lora_transformer_txtfusion_layerwise_blocks_1_attn_wq', 'lora_transformer_txtfusion_refiner_blocks_1_mlp_down'] + if arch == 'anima': + return ['lora_transformer_transformer_blocks_27_attn1_to_q', 'lora_llm_adapter_blocks_5_self_attn_q_proj', 'lora_te_layers_3_mlp_gate_proj'] + if arch == 'zimage': + return ['lora_transformer_layers_29_attention_to_q', 'lora_transformer_noise_refiner_1_attention_to_q'] + if arch == 'sd3': + return ['lora_transformer_transformer_blocks_23_attn_to_q'] + return [] + + +@contextmanager +def block_model(arch, keys=None, **layers): + """mock_model plus a synthetic arch and network_layer_mapping for block classification.""" + from modules import modeldata + real_type = modeldata.get_model_type + modeldata.get_model_type = lambda _pipe: arch + try: + with mock_model(**layers): + shared.sd_model.network_layer_mapping = {k: None for k in (keys or block_fixture_keys(arch))} + lora_blocks.state.update(stamp=None, layout=None) + lora_blocks.state['index'].clear() + lora_blocks.state['vectors'].clear() + lora_blocks.warned.clear() + yield + finally: + modeldata.get_model_type = real_type + lora_blocks.state.update(stamp=None, layout=None) + lora_blocks.state['index'].clear() + lora_blocks.state['vectors'].clear() + lora_blocks.warned.clear() + + +def sd3_spec(n=25, **slots): + """A 25-slot sd3 vector as a spec string with named slot overrides.""" + vals = [1.0] * n + for slot, v in slots.items(): + vals[int(slot[1:])] = v + return ','.join(str(v) for v in vals) + + +def test_block_index_sd_unet_layout(): + with block_model('sd'): + lay = lora_blocks.layout() + assert lay is not None and lay['n'] == 26 and lay['kind'] == 'unet', f'layout={lay}' + cases = { + 'lora_unet_conv_in': 1, + 'lora_unet_down_blocks_0_attentions_0_transformer_blocks_0_attn1_to_q': 2, + 'lora_unet_down_blocks_0_resnets_1_conv1': 3, + 'lora_unet_down_blocks_0_downsamplers_0_conv': 4, + 'lora_unet_down_blocks_2_downsamplers_0_conv': 10, + 'lora_unet_down_blocks_3_resnets_1_conv1': 12, + 'lora_unet_mid_block_attentions_0_transformer_blocks_0_attn2_to_k': 13, + 'lora_unet_up_blocks_0_resnets_0_conv1': 14, + 'lora_unet_up_blocks_1_attentions_2_transformer_blocks_0_ff_net_0_proj': 19, + 'lora_unet_up_blocks_2_upsamplers_0_conv': 22, + 'lora_unet_up_blocks_3_resnets_2_conv1': 25, + 'lora_unet_conv_out': 25, + 'lora_unet_conv_norm_out': 25, + 'lora_unet_time_embedding_linear_1': 0, + 'lora_te_text_model_encoder_layers_0_self_attn_q_proj': 0, + } + for key, expected in cases.items(): + got = lora_blocks.block_index(key) + assert got == expected, f'{key}: got {got} expected {expected}' + return True + + +def test_block_index_sdxl_unet_layout(): + with block_model('sdxl'): + lay = lora_blocks.layout() + assert lay is not None and lay['n'] == 20, f'layout={lay}' + cases = { + 'lora_unet_down_blocks_1_attentions_0_transformer_blocks_3_attn1_to_v': 5, + 'lora_unet_mid_block_attentions_0_transformer_blocks_9_norm3': 10, + 'lora_unet_up_blocks_0_resnets_0_conv1': 11, + 'lora_unet_up_blocks_2_resnets_2_conv1': 19, + 'lora_unet_add_embedding_linear_1': 0, + 'lora_te1_text_model_encoder_layers_0_self_attn_k_proj': 0, + 'lora_te2_text_projection': 0, + } + for key, expected in cases.items(): + got = lora_blocks.block_index(key) + assert got == expected, f'{key}: got {got} expected {expected}' + return True + + +def test_block_index_flux_chains_concatenate(): + with block_model('f1'): + lay = lora_blocks.layout() + assert lay is not None and lay['n'] == 58, f'layout={lay}' # 19 double + 38 single + BASE + cases = { + 'lora_transformer_transformer_blocks_0_attn_to_q': 1, + 'lora_transformer_transformer_blocks_18_ff_net_0_proj': 19, + 'lora_transformer_single_transformer_blocks_0_attn_to_q': 20, + 'lora_transformer_single_transformer_blocks_37_proj_out': 57, + 'lora_transformer_x_embedder': 0, + 'lora_transformer_proj_out': 0, + } + for key, expected in cases.items(): + got = lora_blocks.block_index(key) + assert got == expected, f'{key}: got {got} expected {expected}' + return True + + +def test_block_index_anchoring_krea2_and_chroma(): + with block_model('krea2'): + lay = lora_blocks.layout() + assert lay is not None and lay['n'] == 29, f'layout={lay}' # txtfusion chains stay uncounted + assert lora_blocks.block_index('lora_transformer_blocks_5_attn_wq') == 6 + assert lora_blocks.block_index('lora_transformer_txtfusion_layerwise_blocks_0_attn_wq') == 0 + assert lora_blocks.block_index('lora_transformer_txtfusion_refiner_blocks_1_mlp_down') == 0 + with block_model('chroma'): + assert lora_blocks.block_index('lora_transformer_transformer_blocks_0_attn_to_q') == 1 + assert lora_blocks.block_index('lora_transformer_single_transformer_blocks_0_attn_to_q') == 20 # anchored: not the double chain's slot + assert lora_blocks.block_index('lora_transformer_distilled_guidance_layer_layers_0_linear_1') == 0 + return True + + +def test_block_index_namespace_collisions(): + with block_model('anima'): + assert lora_blocks.block_index('lora_transformer_transformer_blocks_5_attn1_to_q') == 6 + assert lora_blocks.block_index('lora_te_layers_0_self_attn_q_proj') is None, 'anima TE strips to the zimage pattern; the namespace must win' + assert lora_blocks.block_index('lora_llm_adapter_blocks_0_self_attn_q_proj') is None, 'anima llm_adapter strips to the krea2 pattern; the namespace must win' + from types import SimpleNamespace + muted = SimpleNamespace(name='m', block_spec='NONE') + assert lora_blocks.factor('lora_te_layers_0_self_attn_q_proj', muted) == 1.0, 'namespaces outside the vector stay neutral even under an all-zero spec' + return True + + +def test_unet_arithmetic_matches_conversion_map(): + from modules.lora import lora_convert + with block_model('sd'): + n_in = lora_blocks.layout()['n_in'] + checked = 0 + for sd_key, hf_key in lora_convert.make_unet_conversion_map().items(): + if sd_key.startswith('input_blocks'): + expected = 1 + int(sd_key.split('_')[2]) + elif sd_key.startswith('output_blocks'): + expected = 2 + n_in + int(sd_key.split('_')[2]) + elif sd_key.startswith('middle_block'): + expected = 1 + n_in + elif sd_key.startswith('time_embed') or sd_key.startswith('label_emb'): + expected = 0 + elif sd_key.startswith('out_'): + expected = 25 + else: + continue + got = lora_blocks.block_index('lora_unet_' + hf_key) + assert got == expected, f'{sd_key} -> {hf_key}: got {got} expected {expected}' + checked += 1 + assert checked > 60, f'the map cross-check covered only {checked} entries' + return True + + +def test_resolve_preset_case_and_arch_guard(): + from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS + with block_model('sd'): + v = lora_blocks.resolve('grad_v') + assert v is not None and len(v) == 26 and v[0] == 1.0, 'preset BASE must be forced neutral' + assert v[1:] == [float(x) for x in BLOCK_WEIGHTS_PRESETS['GRAD_V'][1:]] + assert lora_blocks.resolve('SDXL_GRAD_V') is None, 'arch-tagged presets must not resolve elsewhere' + with block_model('sdxl'): + v = lora_blocks.resolve('GRAD_V') + assert v is not None and len(v) == 20 and v[0] == 1.0 + assert v[1:] == [float(x) for x in SDXL_BLOCK_WEIGHTS_PRESETS['SDXL_GRAD_V'][1:]], 'the SDXL_ table must serve the unprefixed name' + v = lora_blocks.resolve('RING08_5') + assert v is not None and len(v) == 20 and v[0] == 1.0, 'a 26-slot preset must resample onto the sdxl layout' + return True + + +def test_resolve_dit_resample_drops_base(): + from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS + src = BLOCK_WEIGHTS_PRESETS['GRAD_A'] + with block_model('f1'): + v = lora_blocks.resolve('GRAD_A') + assert v is not None and len(v) == 58 + assert v[0] == 1.0, 'BASE is a unet concept and must not inherit the merge slot' + assert v[1] == float(src[1]) and v[-1] == float(src[-1]), 'resampling must keep the endpoints' + assert min(v[1:]) >= min(src[1:]) - 1e-9 and max(v[1:]) <= max(src[1:]) + 1e-9 + return True + + +def test_resolve_vector_length_policy(): + with block_model('sd'): + full = [round(0.01 * i, 2) for i in range(26)] + v = lora_blocks.resolve(','.join(str(x) for x in full)) + assert v == full, 'a canonical-length vector must pass through verbatim' + v = lora_blocks.resolve(','.join(str(x) for x in full[1:])) + assert v == [1.0] + full[1:], 'a base-less vector must gain a neutral BASE' + v = lora_blocks.resolve(','.join(['0.5'] * 17)) + assert v is not None and len(v) == 26 and v[0] == 0.5 and v[2] == 0.5 and v[13] == 0.5 and v[25] == 0.5, 'the a1111 17-slot layout must expand' + assert v[1] == 1.0 and v[14] == 1.0, 'slots the a1111 layout omits stay neutral' + assert lora_blocks.resolve(','.join(['1'] * 24)) is None, 'an unmatched length must be rejected' + with block_model('sdxl'): + v = lora_blocks.resolve(','.join(['0.25'] * 12)) + assert v is not None and len(v) == 20 and v[0] == 0.25 and v[5] == 0.25 and v[10] == 0.25 and v[16] == 0.25 + assert v[1] == 1.0 and v[7] == 1.0 and v[17] == 1.0 + return True + + +def test_resolve_scalar_and_classic(): + with block_model('sd'): + assert lora_blocks.resolve('0.5') == [0.5] * 26, 'a scalar must broadcast to every slot' + v = lora_blocks.resolve('INS') + assert v[0] == 1.0 and all(x == 1.0 for x in v[1:7]) and all(x == 0.0 for x in v[7:]), f'INS must cover the shallow input half: {v}' + v = lora_blocks.resolve('OUTALL') + assert v[0] == 1.0 and all(x == 0.0 for x in v[1:14]) and all(x == 1.0 for x in v[14:]), f'OUTALL must cover the output side: {v}' + assert lora_blocks.resolve('NONE') == [0.0] * 26 + assert lora_blocks.resolve('DOUBLE') is None, 'chain names need a two-chain arch' + with block_model('f1'): + v = lora_blocks.resolve('DOUBLE') + assert v[0] == 1.0 and all(x == 1.0 for x in v[1:20]) and all(x == 0.0 for x in v[20:]), 'DOUBLE must keep the double chain only' + v = lora_blocks.resolve('SINGLE') + assert all(x == 0.0 for x in v[1:20]) and all(x == 1.0 for x in v[20:]), 'SINGLE must keep the single chain only' + return True + + +def test_bad_value_warns_once_and_ignores(): + from types import SimpleNamespace + with block_model('sd'): + assert lora_blocks.resolve('bogus') is None + assert lora_blocks.resolve('1,2,3') is None + warned_n = len(lora_blocks.warned) + lora_blocks.resolve('bogus') + assert len(lora_blocks.warned) == warned_n, 'a repeated bad value must not warn again' + net = SimpleNamespace(name='b', block_spec='bogus') + assert lora_blocks.factor('lora_unet_conv_in', net) == 1.0, 'an unresolvable spec must leave the plain strength' + return True + + +def test_multiplier_folds_block_weight(): + layer = build_layer('uint4') + layer.network_layer_name = 'lora_transformer_transformer_blocks_3_attn_to_q' + A, B, D = make_delta() + net = make_net('blocky', layer, A, B, te_mult=0.5) + with block_model('sd3', lin=layer): + net.block_spec = sd3_spec(s4=0.5) # block 3 sits in slot 4 + Wdq0 = dq(layer) + activate(net) + rho = rho_of(dq(layer) - Wdq0, D) + assert abs(rho - 0.25) < 0.01, f'expected multiplier 0.5 x block 0.5, rho={rho:.4f}' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact' + return True + + +def test_block_weight_zero_kills_layer_delta(): + layer_a = build_layer('uint4') + layer_a.network_layer_name = 'lora_transformer_transformer_blocks_3_attn_to_q' + layer_b = build_layer('uint4', seed=7) + layer_b.network_layer_name = 'lora_transformer_transformer_blocks_5_attn_to_q' + A1, B1, D1 = make_delta(seed=1) + A2, B2, D2 = make_delta(seed=2) + net = make_net('zeroed', layer_a, A1, B1) + nw = network.NetworkWeights(network_key=layer_b.network_layer_name, sd_key=layer_b.network_layer_name, + w={'lora_up.weight': B2.cpu(), 'lora_down.weight': A2.cpu()}, sd_module=layer_b) + net.modules[layer_b.network_layer_name] = network_lora.NetworkModuleLora(net, nw) + with block_model('sd3', a=layer_a, b=layer_b): + net.block_spec = sd3_spec(s4=0.0) # zero the slot of block 3; block 5 stays at 1 + Wa0, Wb0 = dq(layer_a), dq(layer_b) + activate(net) + rho_a = rho_of(dq(layer_a) - Wa0, D1) + rho_b = rho_of(dq(layer_b) - Wb0, D2) + assert abs(rho_a) < 0.01, f'a zero slot must null the layer delta, rho={rho_a:.4f}' + assert rho_b > 0.99, f'a neutral slot must apply in full, rho={rho_b:.4f}' + activate() + assert torch.equal(dq(layer_a), Wa0) and torch.equal(dq(layer_b), Wb0), 'restore must be bit-exact' + return True + + +def test_signature_suffix_inactive_and_changes(): + from modules.lora import extra_networks_lora + layer = build_layer('uint4') + A, B, _D = make_delta() + net = make_net('siggy', layer, A, B) + l_common.loaded_networks.clear() + l_common.loaded_networks.append(net) + try: + assert lora_blocks.signature() == '', 'no spec must leave the stamp signature untouched' + net.block_spec = 'GRAD_V' + s1 = lora_blocks.signature() + assert s1 == '|lbw=siggy:grad_v', f's1={s1}' + net.block_spec = ' Grad_A ' + assert lora_blocks.signature() == '|lbw=siggy:grad_a', 'the spec must normalize' + en = extra_networks_lora.ExtraNetworkLora() + plain = en.signature(['a'], [1.0], [[1.0] * 3]) + with_spec = en.signature(['a'], [1.0], [[1.0] * 3], ['GRAD_V']) + assert plain == [f'a:1.0:{[1.0] * 3}'], 'legacy signature strings must stay byte-identical without specs' + assert with_spec[0] == plain[0] + ':lbw=grad_v' + finally: + l_common.loaded_networks.clear() + return True + + +def test_factor_cache_invalidates_on_block_weight(): + import tempfile + layer = build_layer('uint4') + layer.network_layer_name = 'lora_transformer_transformer_blocks_3_attn_to_q' + with tempfile.TemporaryDirectory() as tmp: + with host_rank(64), host_cache(10, os.path.join(tmp, 'cache')), block_model('sd3', lin=layer): + net, _D = cache_fixture(tmp, layer) + activate(net) + up_full = layer.svd_up.detach().clone() + activate() + net.block_spec = sd3_spec(s4=0.5) + activate(net) # different block vector: different signature, fresh svd, second entry + assert not torch.equal(layer.svd_up, up_full), 'a block-weight change must produce different factors' + activate() + files = os.listdir(os.path.join(tmp, 'cache')) + assert len(files) == 2, f'two cache entries expected, got {files}' + return True + + +def test_stack_ties_respects_per_net_blocks(): + layer = build_layer('uint4') + layer.network_layer_name = 'lora_transformer_transformer_blocks_3_attn_to_q' + torch.manual_seed(31) + Da = torch.randn(OUT_F, IN_F, device=DEVICE) * 3e-4 + Db = torch.randn(OUT_F, IN_F, device=DEVICE) * 3e-4 + net_a = make_dense_net('tiesa', layer, Da) + net_b = make_dense_net('tiesb', layer, Db) + with host_rank(64), stack_mode('ties', dens=0.5), block_model('sd3', lin=layer): + Wdq0 = dq(layer) + activate(net_a, net_b) + d_both = (dq(layer) - Wdq0).clone() + activate() + net_b.block_spec = 'NONE' + activate(net_a, net_b) + d_muted = (dq(layer) - Wdq0).clone() + activate() + assert not torch.allclose(d_both, d_muted), 'muting one member must change the combined delta' + assert rho_of(d_muted, Db) < 0.1, f'the muted member must not contribute: rho={rho_of(d_muted, Db):.3f}' + assert rho_of(d_muted, Da) > 0.3, f'the live member must survive the trim: rho={rho_of(d_muted, Da):.3f}' + return True + + +def test_pending_promote_updates_block_spec(): + layer = build_layer('uint4') + layer.network_layer_name = 'lora_transformer_transformer_blocks_3_attn_to_q' + A, B, D = make_delta() + net = make_net('promoted', layer, A, B) + with block_model('sd3', lin=layer): + Wdq0 = dq(layer) + net.pending_config = {'te': 1.0, 'unet': [1.0] * 3, 'dyn': None, 'blocks': 'NONE'} + activate(net) + assert net.block_spec == 'NONE', 'network_activate must promote the staged spec' + rho = rho_of(dq(layer) - Wdq0, D) + assert abs(rho) < 0.01, f'the promoted all-zero vector must null the delta, rho={rho:.4f}' + activate() + net.pending_config = {'te': 1.0, 'unet': [1.0] * 3, 'dyn': None, 'blocks': None} + activate(net) + assert net.block_spec is None, 'a spec-less reload must clear the previous spec' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.99, f'without a spec the delta must apply in full, rho={rho:.4f}' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_layout_recomputes_on_mapping_change(): + with block_model('f1'): + assert lora_blocks.layout()['n'] == 58 + assert lora_blocks.block_index('lora_transformer_transformer_blocks_18_attn_to_q') == 19 + shared.sd_model.network_layer_mapping = {'lora_transformer_transformer_blocks_9_attn_to_q': None} # new object: the stamp must miss + assert lora_blocks.layout()['n'] == 11 + assert lora_blocks.block_index('lora_transformer_transformer_blocks_9_attn_to_q') == 10 + assert lora_blocks.block_index('lora_transformer_transformer_blocks_18_attn_to_q') == 0, 'an index past the scanned chain folds to BASE' + return True + + def run_tests(): t0 = time.time() log.warning('=== Erasure law ===') @@ -2380,6 +2766,14 @@ def run_tests(): log.warning('=== Robustness ===') for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: run_test(CAT_ROBUST, fn) + log.warning('=== Block weights ===') + for fn in [test_block_index_sd_unet_layout, test_block_index_sdxl_unet_layout, test_block_index_flux_chains_concatenate, + test_block_index_anchoring_krea2_and_chroma, test_block_index_namespace_collisions, test_unet_arithmetic_matches_conversion_map, + test_resolve_preset_case_and_arch_guard, test_resolve_dit_resample_drops_base, test_resolve_vector_length_policy, + test_resolve_scalar_and_classic, test_bad_value_warns_once_and_ignores, test_multiplier_folds_block_weight, + test_block_weight_zero_kills_layer_delta, test_signature_suffix_inactive_and_changes, test_factor_cache_invalidates_on_block_weight, + test_stack_ties_respects_per_net_blocks, test_pending_promote_updates_block_spec, test_layout_recomputes_on_mapping_change]: + run_test(CAT_BLOCKS, fn) elapsed = time.time() - t0 log.warning('=== Results ===') From b38167314367e4ca9d7d1587cd9ed026a2aff553 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 21:36:07 +0100 Subject: [PATCH 21/23] feat(xyz): lora block weight axis String axis rewriting every lora tag in the prompt: an existing lbw= argument is replaced, None removes it for a clean baseline cell. Choices list the preset names; raw vectors go through csv mode with escaped commas. Long values truncate in the grid legend. --- scripts/xyz/xyz_grid_classes.py | 4 ++++ scripts/xyz/xyz_grid_shared.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/scripts/xyz/xyz_grid_classes.py b/scripts/xyz/xyz_grid_classes.py index 198d93c91..2f1ed59bf 100644 --- a/scripts/xyz/xyz_grid_classes.py +++ b/scripts/xyz/xyz_grid_classes.py @@ -19,6 +19,9 @@ from scripts.xyz.xyz_grid_shared import ( # pylint: disable=no-name-in-module, u list_lora, apply_lora, apply_lora_strength, + list_lora_blocks, + apply_lora_blocks, + format_value_trim, apply_te, apply_guidance, apply_styles, @@ -217,6 +220,7 @@ axis_options = [ AxisOption("[Prompt] Prompt parser", str, apply_setting("prompt_attention"), choices=lambda: ["native", "compel", "xhinker", "a1111", "fixed"]), AxisOption("[Network] LoRA", str, apply_lora, cost=0.5, choices=list_lora), AxisOption("[Network] LoRA strength", float, apply_lora_strength, cost=0.6), + AxisOption("[Network] LoRA block weight", str, apply_lora_blocks, cost=0.6, fmt=format_value_trim, choices=list_lora_blocks), AxisOption("[Network] LoRA stack mode", str, apply_setting("lora_stack_mode"), cost=0.6, choices=lambda: ["sum", "ties", "dare_ties", "dare_linear", "magnitude_prune", "klora", "estlora"]), AxisOption("[Network] LoRA stack density", float, apply_setting("lora_stack_density"), cost=0.6), AxisOption("[Network] LoRA stack ramp", float, apply_setting("lora_stack_alpha"), cost=0.6), diff --git a/scripts/xyz/xyz_grid_shared.py b/scripts/xyz/xyz_grid_shared.py index b90d37848..af6b30b77 100644 --- a/scripts/xyz/xyz_grid_shared.py +++ b/scripts/xyz/xyz_grid_shared.py @@ -260,6 +260,31 @@ def apply_lora_strength(p, x, xs): shared.opts.data['extra_networks_default_multiplier'] = x +def list_lora_blocks(): + from modules.lora import lora_blocks + from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS + return ['None'] + list(lora_blocks.CLASSIC) + list(lora_blocks.CHAIN_NAMES) + sorted(BLOCK_WEIGHTS_PRESETS) + sorted(SDXL_BLOCK_WEIGHTS_PRESETS) + + +re_lora_tag = re.compile(r']+)>') + + +def apply_lora_blocks(p, x, xs): + x = str(x or '').strip() + if ':' in x or '>' in x: + log.error(f'XYZ grid apply LoRA block weight: value="{x}" invalid characters') + return + def rewrite(m): + items = [i for i in m.group(1).split(':') if not i.lower().startswith('lbw=')] + if x and x.lower() != 'none': + items.append(f'lbw={x}') + return '' + p.prompt = re_lora_tag.sub(rewrite, p.prompt) + p.all_prompts = None # a populated list would shadow the edited prompt in processing + p.all_negative_prompts = None + log.debug(f'XYZ grid apply LoRA block weight: "{x}"') + + def apply_te(p, x, xs): shared.opts.data["sd_text_encoder"] = x sd_models.reload_text_encoder() @@ -383,6 +408,13 @@ def format_value_join_list(p, opt, x): return ", ".join(x) +def format_value_trim(p, opt, x): + x = str(x) + if len(x) > 40: + x = x[:37] + '...' # block-weight vectors would flood the grid legend + return f"{opt.label}: {x}" + + def do_nothing(p, x, xs): pass From 999224e91a03a67835fb92a871fdf8ee7963a688 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 20 Aug 2026 22:09:07 +0100 Subject: [PATCH 22/23] test(lora): pin the upstream fixes in the campaign suite The in-place weight installs, the promote-after-deactivate fuse ordering, and the dynamo reset at model unload live in the shared loader code; their regression pins belong in the campaign suite beside the paths they protect. --- test/test-sdnq-lora-factors.py | 117 ++++++++++++++++++++++++++++++++- 1 file changed, 115 insertions(+), 2 deletions(-) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index c0cf7fede..b4308c14c 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -646,6 +646,91 @@ def test_mechanism_flip_restore_pass_strips(): assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'restore pass must strip the factors' assert torch.equal(dq(layer), Wdq0), 'strip must restore bit-exact' assert layer.network_current_names == (), 'stripped layer must be stamped restored' + +def test_apply_restore_preserves_weight_storage(): + """Apply and restore write into the existing parameter storage: kernel selection is + placement-sensitive, so a swapped-in Parameter shifts deterministic outputs bitwise.""" + lin = torch.nn.Linear(IN_F, OUT_F, bias=True, dtype=torch.bfloat16, device=DEVICE) + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.02) + lin.bias.copy_(torch.randn(OUT_F, device=DEVICE) * 0.01) + lin.network_layer_name = 'lora_transformer_storage' + lin.network_current_names = () + W0 = lin.weight.detach().clone() + B0 = lin.bias.detach().clone() + wptr, bptr = lin.weight.data_ptr(), lin.bias.data_ptr() + A, B, _D = make_delta(seed=91, sigma=1e-2) + net = make_net('storage', lin, A, B) + with mock_model(lin=lin): + activate(net) + assert not torch.equal(lin.weight.detach(), W0), 'apply must change the weight' + assert lin.weight.data_ptr() == wptr, 'apply must write into the existing weight storage' + assert lin.bias.data_ptr() == bptr, 'apply must keep the bias storage' + activate() + assert torch.equal(lin.weight.detach(), W0), 'restore must be bit-exact' + assert torch.equal(lin.bias.detach(), B0), 'restore must be bit-exact on bias' + assert lin.weight.data_ptr() == wptr, 'restore must write into the existing weight storage' + assert lin.bias.data_ptr() == bptr, 'restore must write into the existing bias storage' + return True + + +def fuse_fixture(te0): + """A plain bf16 Linear (unquantized, so fuse stays allowed) with one attached net at strength te0.""" + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device=DEVICE) + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.02) + lin.network_layer_name = 'lora_transformer_fusefix' + lin.network_current_names = () + A, B, D = make_delta(seed=17, sigma=1e-2) + net = make_net('fusefix', lin, A, B, te_mult=te0) + return lin, net, D + + +def edit_strength(net, te): + """The production order for a strength edit: network_load stages the new values on the + shared net object, deactivate runs against the applied ones, activate promotes.""" + l_common.previously_loaded_networks[:] = l_common.loaded_networks + net.pending_config = {'te': te, 'unet': [te] * 3, 'dyn': None} + networks.network_deactivate() + networks.network_activate() + + +def test_fuse_promote_applies_new_multiplier(): + """Fuse removal subtracts a recomputed delta, so the multipliers it reads must be the + applied ones: staged values promote only in network_activate, after the removal pass.""" + lin, net, D = fuse_fixture(te0=0.5) + with mock_model(lin=lin): + shared.opts.lora_fuse_native = True + W0 = lin.weight.detach().float().clone() + activate(net) + assert isinstance(getattr(lin, 'network_weights_backup', None), bool), 'fuse mode must not take a tensor backup' + rho0 = rho_of(lin.weight.detach().float() - W0, D) + assert abs(rho0 - 0.5) < 0.05, f'rho={rho0:.3f} expected the initial strength' + edit_strength(net, 1.0) + assert net.te_multiplier == 1.0, 'activate must promote the staged multiplier' + rho1 = rho_of(lin.weight.detach().float() - W0, D) + assert abs(rho1 - 1.0) < 0.05, f'rho={rho1:.3f} expected the edited strength to apply, not the first one' + return True + + +def test_fuse_change_then_remove_restores_pristine(): + """Apply, edit, remove: the final subtraction must use the strength that was applied. + Pins the promote-after-deactivate ordering; a promote that runs before the removal + pass leaves half the delta baked into the weights.""" + lin, net, D = fuse_fixture(te0=0.5) + with mock_model(lin=lin): + shared.opts.lora_fuse_native = True + W0 = lin.weight.detach().float().clone() + activate(net) + edit_strength(net, 1.0) + l_common.previously_loaded_networks[:] = l_common.loaded_networks + l_common.loaded_networks.clear() + networks.network_deactivate() + networks.network_activate() + resid = lin.weight.detach().float() - W0 + rho2 = rho_of(resid, D) + assert abs(rho2) < 0.05, f'rho={rho2:.3f} removal must subtract the strength that was applied' + assert float(resid.abs().max()) < 2e-3, f'max={float(resid.abs().max()):.2e} removal must leave only rounding residue' return True @@ -2275,6 +2360,33 @@ def test_rank_bucket_graph_reuse(): return True +def test_recompile_wall_resets_on_unload(): + import sdnq.common as sdnq_common + if not sdnq_common.use_torch_compile: + return True + import torch._dynamo + import torch._dynamo.config as dcfg + from torch._dynamo.exc import FailOnRecompileLimitHit + old_acc = dcfg.accumulated_recompile_limit + dcfg.accumulated_recompile_limit = 4 + try: + fn = sdnq_common.compile_func(lambda w, s: w.to(torch.float32) * s) + hit = False + for i in range(8): # fresh shapes stand in for model switches: the lifetime counter climbs even when old guards are dead + try: + fn(torch.randint(0, 255, (32 + 16 * i, 8), dtype=torch.uint8, device=DEVICE), torch.rand(32 + 16 * i, 1, device=DEVICE)) + except FailOnRecompileLimitHit: + hit = True + break + assert hit, 'the lowered lifetime wall must trip on fullgraph recompiles' + sdnq_common.reset_compile_caches() # the unload-seam hook: counters and dead graphs cleared + fn(torch.randint(0, 255, (1024, 8), dtype=torch.uint8, device=DEVICE), torch.rand(1024, 1, device=DEVICE)) + finally: + dcfg.accumulated_recompile_limit = old_acc + torch._dynamo.reset() # leave no wall residue for later tests + return True + + CAT_ROBUST = category('robustness') @@ -2724,7 +2836,8 @@ def run_tests(): log.warning('=== Set transitions ===') for fn in [test_mixed_family_transition_restores_base, test_partial_coverage_layers_stay_independent, test_mechanism_gate_declines_candidates, test_requantize_option_routes_to_legacy_path, - test_mechanism_flip_strips_attached_factors, test_mechanism_flip_restore_pass_strips]: + test_mechanism_flip_strips_attached_factors, test_mechanism_flip_restore_pass_strips, + test_apply_restore_preserves_weight_storage, test_fuse_promote_applies_new_multiplier, test_fuse_change_then_remove_restores_pristine]: run_test(CAT_TRANS, fn) log.warning('=== Hosting ===') for fn in [test_hosted_low_rank_delta_is_kept, test_hosted_dense_delta_beats_requant, test_hosted_skips_int8, @@ -2761,7 +2874,7 @@ def run_tests(): test_select_reset_reports_timing, test_select_weight_flip_calcs_on_accelerator]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===') - for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: + for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse, test_recompile_wall_resets_on_unload]: run_test(CAT_COMPILE, fn) log.warning('=== Robustness ===') for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: From 5cfa07fb5e57064303d5ca23e58919fd70604b54 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 20 Aug 2026 22:44:02 +0100 Subject: [PATCH 23/23] test(lora): extend apply-method coverage to select riding The mechanism gate tests assert select_candidate declines under requantize, and the apply-method hint names the cache option among those the requantize choice disables. --- modules/lora/networks.py | 2 +- test/test-sdnq-lora-factors.py | 185 +++++++++++++++++---------------- 2 files changed, 95 insertions(+), 92 deletions(-) diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 479ea998a..45afb8480 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -92,7 +92,7 @@ def network_activate(include=None, exclude=None): refused = 0 with devices.inference_context(), pbar: wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () - stack_sig = lora_stack.signature() + lora_blocks.signature() + lora_sdnq.signature() # tracked beside network_current_names so settings-only stack, block-weight or mechanism changes re-apply + stack_sig = lora_stack.signature() + lora_blocks.signature() + lora_sdnq.signature() # tracked beside network_current_names so stack-setting, block-weight and mechanism changes re-apply select_active = len(l.loaded_networks) > 0 and lora_stack.active_select(len(l.loaded_networks)) # restore-only walks have nothing to stack; the count warning would fire on every network-free generation applied_layers.clear() lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index b4308c14c..af44cb1b0 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -558,95 +558,6 @@ def test_partial_coverage_layers_stay_independent(): return True -@contextmanager -def apply_method(value): - old = getattr(shared.opts, 'lora_sdnq_apply', 'exact') - shared.opts.lora_sdnq_apply = value - try: - yield - finally: - shared.opts.lora_sdnq_apply = old - - -def test_mechanism_gate_declines_candidates(): - """The requantize option must gate every svd-channel entry point and flip the apply-stamp token.""" - layer = build_layer('uint4') - A, B, _D = make_delta() - l_common.loaded_networks.clear() - l_common.loaded_networks.append(make_net('one', layer, A, B)) - wanted = (('one', 1.0, 1.0, None),) - try: - assert lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) - with host_rank(64): - assert lora_sdnq.host_candidate(layer, layer.network_layer_name, wanted) - assert lora_sdnq.signature() == '' - with apply_method('requantize'): - assert not lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) - with host_rank(64): - assert not lora_sdnq.host_candidate(layer, layer.network_layer_name, wanted) - assert lora_sdnq.signature() == '|quant=requantize' - finally: - l_common.loaded_networks.clear() - return True - - -def test_requantize_option_routes_to_legacy_path(): - """With the option set, a factorable set must take the classic backup-and-requantize path end to end.""" - layer = build_layer('uint4') - A, B, _D = make_delta(sigma=3e-3) - net = make_net('one', layer, A, B) - with apply_method('requantize'), mock_model(lin=layer): - shared.opts.lora_fuse_native = False # a real quantized model forces backup mode; the mock carries no quantization config - Wdq0 = dq(layer) - activate(net) - assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'legacy path must not touch the svd channel' - assert layer.svd_up is None, 'legacy path must leave the channel empty' - assert isinstance(layer.network_weights_backup, torch.Tensor), 'legacy path must take a tensor backup' - assert not torch.equal(dq(layer), Wdq0), 'legacy path must requantize the weights' - activate() - assert torch.equal(dq(layer), Wdq0), 'legacy restore must be bit-exact from backup' - return True - - -def test_mechanism_flip_strips_attached_factors(): - """Flipping to requantize with factors attached must strip them before the weight path takes the layer; flipping back must re-enter the factor path.""" - layer = build_layer('uint4') - A, B, _D = make_delta(sigma=3e-3) - net = make_net('one', layer, A, B) - with mock_model(lin=layer): - shared.opts.lora_fuse_native = False # a real quantized model forces backup mode; the mock carries no quantization config - Wdq0 = dq(layer) - activate(net) - assert hasattr(layer, 'sdnq_lora_svd_stash'), 'default mechanism must take the factor path' - E_exact = dq(layer) - Wdq0 - with apply_method('requantize'): - activate(net) # same set; the mechanism token in the apply stamp must force re-processing - assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'flip must strip the attached factors' - assert layer.svd_up is None, 'stripped channel must be empty, or the requantized delta double-applies' - assert isinstance(layer.network_weights_backup, torch.Tensor), 'flipped layer must continue on the backup path' - activate(net) # flip back within the same loaded set - assert hasattr(layer, 'sdnq_lora_svd_stash'), 'flip back must re-enter the factor path' - assert torch.equal(dq(layer) - Wdq0, E_exact), 'exact re-apply must restore the base from backup before attaching' - activate() - assert torch.equal(dq(layer), Wdq0), 'unload must return bit-exact pristine' - return True - - -def test_mechanism_flip_restore_pass_strips(): - """A restore-only pass under the requantize option must still drop attached factors.""" - layer = build_layer('uint4') - A, B, _D = make_delta(sigma=3e-3) - net = make_net('one', layer, A, B) - with mock_model(lin=layer): - Wdq0 = dq(layer) - activate(net) - assert hasattr(layer, 'sdnq_lora_svd_stash') - with apply_method('requantize'): - activate() # unload with the gate closed: the fallthrough strip is the only removal route - assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'restore pass must strip the factors' - assert torch.equal(dq(layer), Wdq0), 'strip must restore bit-exact' - assert layer.network_current_names == (), 'stripped layer must be stamped restored' - def test_apply_restore_preserves_weight_storage(): """Apply and restore write into the existing parameter storage: kernel selection is placement-sensitive, so a swapped-in Parameter shifts deterministic outputs bitwise.""" @@ -734,6 +645,98 @@ def test_fuse_change_then_remove_restores_pristine(): return True +@contextmanager +def apply_method(value): + old = getattr(shared.opts, 'lora_sdnq_apply', 'exact') + shared.opts.lora_sdnq_apply = value + try: + yield + finally: + shared.opts.lora_sdnq_apply = old + + +def test_mechanism_gate_declines_candidates(): + """The requantize option must gate every svd-channel entry point and flip the apply-stamp token.""" + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('one', layer, A, B)) + wanted = (('one', 1.0, 1.0, None),) + try: + assert lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) + with host_rank(64): + assert lora_sdnq.select_candidate(layer, layer.network_layer_name, wanted) + assert lora_sdnq.signature() == '' + with apply_method('requantize'): + assert not lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) + with host_rank(64): + assert not lora_sdnq.select_candidate(layer, layer.network_layer_name, wanted) + assert not lora_sdnq.host_candidate(layer, layer.network_layer_name, wanted) + assert lora_sdnq.signature() == '|quant=requantize' + finally: + l_common.loaded_networks.clear() + return True + + +def test_requantize_option_routes_to_legacy_path(): + """With the option set, a factorable set must take the classic backup-and-requantize path end to end.""" + layer = build_layer('uint4') + A, B, _D = make_delta(sigma=3e-3) + net = make_net('one', layer, A, B) + with apply_method('requantize'), mock_model(lin=layer): + shared.opts.lora_fuse_native = False # a real quantized model forces backup mode; the mock carries no quantization config + Wdq0 = dq(layer) + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'legacy path must not touch the svd channel' + assert layer.svd_up is None, 'legacy path must leave the channel empty' + assert isinstance(layer.network_weights_backup, torch.Tensor), 'legacy path must take a tensor backup' + assert not torch.equal(dq(layer), Wdq0), 'legacy path must requantize the weights' + activate() + assert torch.equal(dq(layer), Wdq0), 'legacy restore must be bit-exact from backup' + return True + + +def test_mechanism_flip_strips_attached_factors(): + """Flipping to requantize with factors attached must strip them before the weight path takes the layer; flipping back must re-enter the factor path.""" + layer = build_layer('uint4') + A, B, _D = make_delta(sigma=3e-3) + net = make_net('one', layer, A, B) + with mock_model(lin=layer): + shared.opts.lora_fuse_native = False # a real quantized model forces backup mode; the mock carries no quantization config + Wdq0 = dq(layer) + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'default mechanism must take the factor path' + E_exact = dq(layer) - Wdq0 + with apply_method('requantize'): + activate(net) # same set; the mechanism token in the apply stamp must force re-processing + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'flip must strip the attached factors' + assert layer.svd_up is None, 'stripped channel must be empty, or the requantized delta double-applies' + assert isinstance(layer.network_weights_backup, torch.Tensor), 'flipped layer must continue on the backup path' + activate(net) # flip back within the same loaded set + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'flip back must re-enter the factor path' + assert torch.equal(dq(layer) - Wdq0, E_exact), 'exact re-apply must restore the base from backup before attaching' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must return bit-exact pristine' + return True + + +def test_mechanism_flip_restore_pass_strips(): + """A restore-only pass under the requantize option must still drop attached factors.""" + layer = build_layer('uint4') + A, B, _D = make_delta(sigma=3e-3) + net = make_net('one', layer, A, B) + with mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash') + with apply_method('requantize'): + activate() # unload with the gate closed: the fallthrough strip is the only removal route + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'restore pass must strip the factors' + assert torch.equal(dq(layer), Wdq0), 'strip must restore bit-exact' + assert layer.network_current_names == (), 'stripped layer must be stamped restored' + return True + + CAT_HOST = category('hosting') @@ -2835,9 +2838,9 @@ def run_tests(): run_test(CAT_E2E, fn) log.warning('=== Set transitions ===') for fn in [test_mixed_family_transition_restores_base, test_partial_coverage_layers_stay_independent, + test_apply_restore_preserves_weight_storage, test_fuse_promote_applies_new_multiplier, test_fuse_change_then_remove_restores_pristine, test_mechanism_gate_declines_candidates, test_requantize_option_routes_to_legacy_path, - test_mechanism_flip_strips_attached_factors, test_mechanism_flip_restore_pass_strips, - test_apply_restore_preserves_weight_storage, test_fuse_promote_applies_new_multiplier, test_fuse_change_then_remove_restores_pristine]: + test_mechanism_flip_strips_attached_factors, test_mechanism_flip_restore_pass_strips]: run_test(CAT_TRANS, fn) log.warning('=== Hosting ===') for fn in [test_hosted_low_rank_delta_is_kept, test_hosted_dense_delta_beats_requant, test_hosted_skips_int8,