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/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index d0a6ffe85..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,12 +176,13 @@ 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 - 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" @@ -220,11 +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 @@ -251,12 +259,12 @@ 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') 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]}') @@ -269,7 +277,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={"fuse" if lora_overrides.fuse_native() else "backup"} 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): diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 3cb8688c0..11b91d9e9 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 @@ -67,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) @@ -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 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) @@ -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,17 @@ 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() + 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_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 2b00bbcc7..1da7106c6 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -51,12 +51,15 @@ 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': [], } + 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) @@ -64,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 @@ -155,6 +162,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_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/lora_overrides.py b/modules/lora/lora_overrides.py index eb629896e..305d1e0e6 100644 --- a/modules/lora/lora_overrides.py +++ b/modules/lora/lora_overrides.py @@ -112,6 +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.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_sdnq.py b/modules/lora/lora_sdnq.py index d95849527..35635b227 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 @@ -55,6 +55,8 @@ from modules.logger import log 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 @@ -137,6 +139,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 @@ -171,6 +176,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 @@ -207,15 +213,28 @@ 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 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] @@ -237,10 +256,11 @@ 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): - """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: @@ -249,11 +269,20 @@ 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 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 any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks) + 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 def apply_cached(self, network_layer_name, wanted_names): @@ -277,15 +306,17 @@ 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 len(members) == 0 and self.svd_up is None: + 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 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 @@ -328,21 +359,23 @@ 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 # 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 @@ -427,6 +460,104 @@ 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. + + 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)) + 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 + 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: @@ -437,6 +568,12 @@ 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(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] @@ -445,7 +582,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() @@ -457,7 +594,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 new file mode 100644 index 000000000..85818c8b9 --- /dev/null +++ b/modules/lora/lora_stack.py @@ -0,0 +1,405 @@ +"""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 time +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, 'gamma_e': 1.0, 'total_steps': 0, 'finalized': False, 'reported': None} +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 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: + log.warning(f'Network stack: mode={m} networks={n_loaded} required=2 fallback=sum') + return False + if select_blocked(): + log.warning(f'Network stack: mode={m} compile=model fallback=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_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. + + 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, 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 + 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], 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: + state['finalized'] = False + + +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['gamma_e'] = 1.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): + """Record a select-mode layer for schedule finalization. + + 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 = {'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()) + 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 + # 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 * state['gamma_e']: + return step + 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, silent=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) + # 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'] = {} + 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 + style_first += initial + 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'] + 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}') + # 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): + """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'] or int(total_steps) <= 0: + 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 + 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): + 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_seg, drop_seg = ((t0, t1), (s0, s1)) if winner == 1 else ((s0, s1), (t0, t1)) + stash = entry['stash'][winner] + 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 + 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') + 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 + if net_module is None: + return + weight = getattr(module, 'weight', None) + 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 + 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/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 9e3b2e7a0..45afb8480 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -2,9 +2,11 @@ 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 +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 @@ -54,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: @@ -62,7 +65,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 @@ -89,10 +92,13 @@ 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_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 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 () @@ -100,13 +106,62 @@ 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 + 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 + if select_active and component_wanted and not network_layer_name.startswith('lora_te'): + 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) + 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): + 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 + 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, 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) + 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): @@ -152,6 +207,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: @@ -186,6 +242,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') @@ -196,6 +253,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 = [] @@ -213,7 +279,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/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_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/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/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 diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index b68f449d0..3d6f02a08 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -708,7 +708,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"), @@ -718,6 +717,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(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), "lora_add_hashes_to_infotext": OptionInfo(False, "LoRA add hash info to metadata"), })) diff --git a/scripts/xyz/xyz_grid_classes.py b/scripts/xyz/xyz_grid_classes.py index 5ce14b422..f8ff5d05c 100644 --- a/scripts/xyz/xyz_grid_classes.py +++ b/scripts/xyz/xyz_grid_classes.py @@ -25,6 +25,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, @@ -110,6 +113,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 @@ -143,6 +150,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 @@ -156,6 +167,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 @@ -214,6 +229,11 @@ 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), + 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/scripts/xyz/xyz_grid_shared.py b/scripts/xyz/xyz_grid_shared.py index f84696f41..bb6deabdb 100644 --- a/scripts/xyz/xyz_grid_shared.py +++ b/scripts/xyz/xyz_grid_shared.py @@ -329,6 +329,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() @@ -452,6 +477,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 diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 3b745ad11..af44cb1b0 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -76,12 +76,13 @@ 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_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 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] = {} @@ -557,6 +558,93 @@ def test_partial_coverage_layers_stay_independent(): return True +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 + + @contextmanager def apply_method(value): old = getattr(shared.opts, 'lora_sdnq_apply', 'exact') @@ -577,11 +665,12 @@ def test_mechanism_gate_declines_candidates(): 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.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: @@ -871,6 +960,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') @@ -1403,6 +1505,768 @@ 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_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) + 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_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(): + # 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, layer=layer) + with mock_model(lin=layer), select_mode('estlora', alpha=1.0, disc=0.5): + Wdq0 = dq(layer) + activate(n1, n2) + 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 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 + + +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_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' + 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) + 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 + + +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 + + +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 + + +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') @@ -1499,6 +2363,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') @@ -1545,6 +2436,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 ===') @@ -1561,6 +2838,7 @@ 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]: run_test(CAT_TRANS, fn) @@ -1568,7 +2846,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 ===') @@ -1579,15 +2857,39 @@ 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_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 ===') + 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, + 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_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]: 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 ===') diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 04c7ae34c..07665ddb2 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -861,13 +861,16 @@ {"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"}, - {"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 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; 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.

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"}, {"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"},