diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index de167af24..bc07d61d0 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -73,6 +73,8 @@ def infotext(p): names = [i.name for i in l.loaded_networks] if len(names) > 0: p.extra_generation_params["LoRA networks"] = ", ".join(names) + if networks.refused_writes > 0: # the model took only part of what was loaded, so the names above do not describe what generated the image + p.extra_generation_params["LoRA refused"] = networks.refused_writes if shared.opts.lora_add_hashes_to_infotext: network_hashes = [] for item in l.loaded_networks: diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index b1225c885..ba44cd27e 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -264,15 +264,21 @@ def new_network(name, network_on_disk): def finalize_network(net, name, family, lora_scale, t0, unmapped=0, mismatch=0, skipped=0): """Emit the standard debug log line and return the populated network (or ``None``). - Returns ``None`` when no modules were bound. Logs at debug only; loader - callers can surface higher-level outcomes at info if needed. + Returns ``None`` when no modules were bound. Records ``mismatch`` on the + network so :func:`try_load_chain` can refuse the file as a whole. """ + net.mismatch = mismatch if len(net.modules) == 0: - if unmapped or mismatch or skipped: - log.debug( + if mismatch: + log.error( f'Network load: type={family} name="{name}" native no-match' f' unmapped={unmapped} mismatch={mismatch} skipped={skipped}' ) + elif unmapped or skipped: + log.debug( + f'Network load: type={family} name="{name}" native no-match' + f' unmapped={unmapped} skipped={skipped}' + ) return None log.debug( f'Network load: type={family} name="{name}" native modules={len(net.modules)}' @@ -328,6 +334,19 @@ def lokr_kron_shape(w): return r1 * r2, c1 * c2_flat +def bias_delta_fits(sd_module, bias_w: torch.Tensor) -> bool: + """Bias-delta sanity check against the live module bias. + + A module with no bias is not a mismatch: whole architectures are built + ``bias=False`` (flux2 has not one biased module), so a stray delta there is + one unappliable key and the apply pass counts it refused. + """ + bias = getattr(sd_module, "bias", None) + if bias is None: + return True + return tuple(bias_w.shape) == tuple(bias.shape) + + def lokr_shapes_match(sd_module, kron_shape, chunk: ChunkSpec | None) -> bool: """Kron-vs-module dim check, honoring SDNQ original shapes and chunk rows. @@ -622,6 +641,15 @@ def try_load_lora(name, network_on_disk, lora_scale, *, mismatch += 1 continue + if "diff_b" in target_w and not bias_delta_fits(sd_module, target_w["diff_b"]): + log.warning( + f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key}' + f' bias={tuple(target_w["diff_b"].shape)} module={tuple(sd_module.bias.shape)}' + f' bias shape mismatch' + ) + mismatch += 1 + continue + nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=target_w, sd_module=sd_module) net.modules[network_key] = network_lora.NetworkModuleLora(net, nw) @@ -976,6 +1004,7 @@ def try_load_norm(name, network_on_disk, lora_scale, *, ) unmapped = 0 + mismatch = 0 for (prefix, base), w in groups.items(): if "w_norm" not in w: continue @@ -992,12 +1021,20 @@ def try_load_norm(name, network_on_disk, lora_scale, *, if sd_module is None: unmapped += 1 continue + if "b_norm" in w and not bias_delta_fits(sd_module, w["b_norm"]): + log.warning( + f'Network load: type=Norm name="{name}" arch={arch_name} key={network_key}' + f' bias={tuple(w["b_norm"].shape)} module={tuple(sd_module.bias.shape)}' + f' bias shape mismatch' + ) + mismatch += 1 + continue if not getattr(sd_module, "network_layer_name", None): sd_module.network_layer_name = network_key nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) net.modules[network_key] = network_norm.NetworkModuleNorm(net, nw) - return finalize_network(net, name, "Norm", lora_scale, t0, unmapped=unmapped) + return finalize_network(net, name, "Norm", lora_scale, t0, unmapped=unmapped, mismatch=mismatch) def try_load_full(name, network_on_disk, lora_scale, *, @@ -1029,6 +1066,7 @@ def try_load_full(name, network_on_disk, lora_scale, *, unmapped = 0 skipped = 0 + mismatch = 0 for (prefix, base), w in groups.items(): if "diff" not in w: continue @@ -1044,6 +1082,14 @@ def try_load_full(name, network_on_disk, lora_scale, *, if sd_module is None: unmapped += 1 continue + if "diff_b" in w and not bias_delta_fits(sd_module, w["diff_b"]): + log.warning( + f'Network load: type=Full name="{name}" arch={arch_name} key={network_key}' + f' bias={tuple(w["diff_b"].shape)} module={tuple(sd_module.bias.shape)}' + f' bias shape mismatch' + ) + mismatch += 1 + continue # Loader-local stamping, same as try_load_norm: a full-weight extraction carries the # norm weights too, and assign_network_names_to_compvis_modules puts norms in the # mapping but never stamps network_layer_name on them, which is what the apply pass @@ -1053,7 +1099,7 @@ def try_load_full(name, network_on_disk, lora_scale, *, nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) net.modules[network_key] = network_full.NetworkModuleFull(net, nw) - return finalize_network(net, name, "Full", lora_scale, t0, unmapped=unmapped, skipped=skipped) + return finalize_network(net, name, "Full", lora_scale, t0, unmapped=unmapped, mismatch=mismatch, skipped=skipped) # === Per-arch umbrella === @@ -1070,13 +1116,18 @@ def try_load_chain(name, network_on_disk, lora_scale, family_loaders): """ sd_models_utils.state_dict_cache.enable() net = None + mismatch = 0 for try_fn in family_loaders: sub = try_fn(name, network_on_disk, lora_scale) if sub is None: continue + mismatch += getattr(sub, 'mismatch', 0) if net is None: net = sub else: net.modules.update(sub.modules) sd_models_utils.state_dict_cache.disable() + if net is not None and mismatch > 0: # applying only the layers that fit leaves the model in a state nothing was trained for + log.error(f'Network load: type=LoRA name="{name}" modules={len(net.modules)} mismatch={mismatch} shapes do not match the loaded model') + return None return net diff --git a/modules/lora/network.py b/modules/lora/network.py index 213f979a6..b16060563 100644 --- a/modules/lora/network.py +++ b/modules/lora/network.py @@ -150,6 +150,7 @@ class Network: # LoraModule self.dyn_dim = None 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 self.bundle_embeddings = {} self.mtime = None self.mentioned_name = None diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 08460dafa..31945c27d 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -9,6 +9,7 @@ from modules.logger import log, console applied_layers: list[str] = [] +refused_writes: int = 0 # deltas the modules would not take on the last activate pass; infotext reports the network as partial native_active: bool = False default_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'text_encoder_4', 'unet', 'transformer', 'transformer_2', 'llm_adapter'] @@ -127,8 +128,9 @@ def network_activate(include=None, exclude=None): if task is not None and len(applied_layers) == 0: pbar.remove_task(task) # hide progress bar for no action - global native_active # pylint: disable=global-statement + global native_active, refused_writes # pylint: disable=global-statement native_active = len(l.loaded_networks) > 0 + refused_writes = refused 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') diff --git a/test/test-krea2-native-adapters.py b/test/test-krea2-native-adapters.py index 29efe3507..9b80b0879 100644 --- a/test/test-krea2-native-adapters.py +++ b/test/test-krea2-native-adapters.py @@ -443,6 +443,49 @@ def test_lora_official_diffusers_renamed(): return True +def test_lora_bias_delta_binds(): + """A diff_b sized to the module bias rides along with the weight LoRA.""" + sd = lora_pair('diffusion_model.first', 'first') + sd['diffusion_model.first.diff_b'] = torch.randn(ckpt_shape('first')[0]) + net = _load_via(K.try_load_lora, sd) + assert net is not None and 'lora_transformer_first' in net.modules, f'got {set(net.modules) if net else None}' + assert net.mismatch == 0, f'mismatch={net.mismatch}' + return True + + +def test_lora_bias_delta_wrong_shape_rejected(): + """A diff_b that does not fit the module bias is a mismatch, not an apply-time surprise.""" + sd = lora_pair('diffusion_model.first', 'first') + sd['diffusion_model.first.diff_b'] = torch.randn(ckpt_shape('first')[0] + 3) + net = _load_via(K.try_load_lora, sd) + assert net is None, f'expected rejection, got {net.modules}' + return True + + +def test_lora_bias_delta_on_biasless_module_binds(): + """A bias delta aimed at a module built without one is not a mismatch. + + The krea2 blocks are ``bias=False`` and whole arches (flux2) carry no bias + at all, so a stray delta there is one unappliable key rather than the wrong + file. It binds, and the apply pass counts it refused. + """ + sd = lora_pair('diffusion_model.blocks.0.attn.wq', 'blocks.0.attn.wq') + sd['diffusion_model.blocks.0.attn.wq.diff_b'] = torch.randn(ckpt_shape('blocks.0.attn.wq')[0]) + net = _load_via(K.try_load_lora, sd) + assert net is not None and net.mismatch == 0, f'got {net.mismatch if net else None}' + return True + + +def test_chain_refuses_whole_network_on_mismatch(): + """One bad delta refuses the file rather than applying the layers that fit.""" + sd = lora_pair('diffusion_model.blocks.0.attn.wq', 'blocks.0.attn.wq') + sd.update(lora_pair('diffusion_model.blocks.0.mlp.up', 'blocks.0.mlp.up')) + sd['diffusion_model.blocks.0.mlp.up.lora_A.weight'] = torch.randn(RANK, ckpt_shape('blocks.0.mlp.up')[1] + 8) + net = _load_via(K.try_load, sd) + assert net is None, f'expected refusal, got {set(net.modules)}' + return True + + def test_lora_bare_diffusers_renamed(): """Bare diffusers key (save_lora_adapter output) renames and binds.""" sd = lora_pair('transformer_blocks.1.ff.up', 'blocks.1.mlp.up') @@ -551,14 +594,18 @@ def test_oft_diffusers_renamed(): def test_full_diff_chain(): - """Full-diff extraction loads through the try_load chain and yields finite updown.""" - out, inp = ckpt_shape('blocks.0.attn.wq') + """Full-diff extraction loads through the try_load chain and yields finite updown. + + Targets ``img_in``/``first`` rather than a block attention leaf: the blocks are + built ``bias=False``, so a diff_b aimed at one is a delta with nothing to land on. + """ + out, inp = ckpt_shape('first') sd = { - 'transformer.transformer_blocks.0.attn.to_q.diff': torch.randn(out, inp), - 'transformer.transformer_blocks.0.attn.to_q.diff_b': torch.randn(out), + 'transformer.img_in.diff': torch.randn(out, inp), + 'transformer.img_in.diff_b': torch.randn(out), } net = _load_via(K.try_load, sd) - assert net is not None and 'lora_transformer_blocks_0_attn_wq' in net.modules, f'got {set(net.modules) if net else None}' + assert net is not None and 'lora_transformer_first' in net.modules, f'got {set(net.modules) if net else None}' mod = next(iter(net.modules.values())) updown, ex_bias = mod.calc_updown(mod.sd_module.weight) assert tuple(updown.shape) == (out, inp) and torch.isfinite(updown).all() @@ -645,6 +692,10 @@ def run_tests(): log.warning('=== Loaders ===') for fn in [ test_lora_official_diffusers_renamed, + test_lora_bias_delta_binds, + test_lora_bias_delta_wrong_shape_rejected, + test_lora_bias_delta_on_biasless_module_binds, + test_chain_refuses_whole_network_on_mismatch, test_lora_bare_diffusers_renamed, test_lora_comfy_checkpoint_verbatim, test_lora_kohya_checkpoint, diff --git a/test/test-lora-apply.py b/test/test-lora-apply.py index c8a0dac60..4446d0010 100644 --- a/test/test-lora-apply.py +++ b/test/test-lora-apply.py @@ -101,9 +101,13 @@ def make_linear(out_features: int, in_features: int, seed: int = 0): def stamp_fuse(module): - """Mark the module as network_backup_weights leaves it in fuse mode: no tensor backup.""" + """Mark the module as network_backup_weights leaves it in fuse mode: no tensor backup. + + The bias flag is only set when the module has one, same as the loader does. + """ module.network_weights_backup = True - module.network_bias_backup = True + if getattr(module, 'bias', None) is not None: + module.network_bias_backup = True def stamp_backup(module, weight, bias): @@ -187,6 +191,23 @@ def test_fuse_mismatched_bias_delta_is_refused(): return True +def test_fuse_bias_delta_without_a_bias_is_refused(): + """A delta aimed at a bias the module does not have is counted, not silently dropped. + + The loader lets this through on purpose: whole architectures are built + bias=False, so a stray diff_b is one unappliable key rather than the wrong + file. The apply pass is where it has to become visible. + """ + module = torch.nn.Linear(8, 32, bias=False) + w0 = module.weight.detach().clone() + stamp_fuse(module) + written = network_apply_direct(module, torch.full_like(w0, 0.5), torch.full((32,), 0.25), device=CPU) + assert written == (True, False), f'reported {written}' + assert module.bias is None, 'a bias appeared on a module built without one' + assert_close(module.weight.detach(), w0 + 0.5, 'weight') + return True + + def test_fuse_mismatched_weight_delta_is_refused(): """A weight delta that does not fit is dropped while the bias delta still lands.""" module, w0, b0 = make_linear(32, 8) @@ -248,6 +269,7 @@ def run_tests(): test_fuse_deactivate_restores, test_fuse_mismatched_bias_delta_is_refused, test_fuse_mismatched_weight_delta_is_refused, + test_fuse_bias_delta_without_a_bias_is_refused, ]: run_test(CAT_FUSE, fn)