From 6937ea803fb51c357dee3bf81e3eb1bdea88ee61 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 23:34:41 +0100 Subject: [PATCH] refactor(ernie): migrate to generic native_loader Replaces ernie's four family loaders with thin wrappers binding native_loader's generics to ernie's prefix tuples and resolve_targets. ErnieImageAttention has fully split to_q / to_k / to_v with no fused QKV and ErnieImageFeedForward has three separate Linear modules, so resolve_targets is a straight passthrough across every recognized prefix. BARE_DIFFUSERS_PREFIXES covers layers., adaLN_modulation., final_norm., final_linear. for bare-diffusers exports (e.g. via save_lora_adapter). parse_key returns (prefix_used, base, suffix) instead of the old (network_key, suffix); parse test updated. --- pipelines/ernie/ernie_lora.py | 337 ++++++++--------------------- test/test-ernie-native-adapters.py | 14 +- 2 files changed, 95 insertions(+), 256 deletions(-) diff --git a/pipelines/ernie/ernie_lora.py b/pipelines/ernie/ernie_lora.py index 3401a4bec..c7ad8d497 100644 --- a/pipelines/ernie/ernie_lora.py +++ b/pipelines/ernie/ernie_lora.py @@ -1,273 +1,110 @@ """ERNIE-Image native adapter loader. Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'`` -(``lora_force_diffusers`` off and ``ernieimage`` in ``allow_native``). Reads the -safetensors directly and writes into sdnext's existing -``network_layer_mapping``, returning a ``Network`` populated with -``NetworkModule*`` entries that ``network_activate`` will apply. If the -setting is on, the diffusers PEFT path handles the file instead. +(``lora_force_diffusers`` off and ``ernieimage`` in ``allow_native``). -Entry points, one per family: +Entry points, one per family: :func:`try_load_lora` (plus DoRA), +:func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`. -- LoRA (+ DoRA) via :func:`try_load_lora` -- LoKR via :func:`try_load_lokr` -- LoHA via :func:`try_load_loha` -- OFT via :func:`try_load_oft` +Recognized key prefixes: ``diffusion_model.``, ``transformer.``, +``lora_unet_``, plus bare diffusers paths (``layers.``, ``adaLN_modulation.``, +``final_norm.``, ``final_linear.``). -Recognized key prefixes for every family: ``diffusion_model.``, -``transformer.``, ``lora_unet_``, or bare. Diffusers-PEFT ``lora_A``/``lora_B`` -are normalized to ``lora_down``/``lora_up``. - -The ERNIE-Image transformer has separate ``self_attention.to_q``/``to_k``/ -``to_v``/``to_out.0`` linear modules (no fused QKV layout), so no -chunk/split machinery is needed and all four families are supported uniformly. +``ErnieImageAttention`` has fully split ``to_q`` / ``to_k`` / ``to_v`` Linear +modules (no fused QKV) and ``ErnieImageFeedForward`` exposes ``gate_proj``, +``up_proj``, ``linear_fc2`` separately. resolve_targets is therefore a +straight passthrough; no chunking, no renames, no dispatch table. """ -import os -import time -import torch -from modules import shared, sd_models -from modules.logger import log -from modules.lora import network, network_lora, network_lokr, network_hada, network_oft, lora_convert -from modules.lora import lora_common as l +from modules.lora import native_loader -KNOWN_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_") +# === Arch-specific prefix configuration === -# Every family also picks up the universal optional keys -# (alpha, scale, bias, dora_scale) via base NetworkModule.__init__. -LORA_SUFFIXES = ( - ".lora_down.weight", ".lora_up.weight", - ".lora_A.weight", ".lora_B.weight", - ".alpha", ".dora_scale", ".bias", ".scale", -) -LOKR_SUFFIXES = ( - ".lokr_w1", ".lokr_w2", - ".lokr_w1_a", ".lokr_w1_b", - ".lokr_w2_a", ".lokr_w2_b", - ".lokr_t2", - ".alpha", ".dora_scale", ".bias", ".scale", -) -LOHA_SUFFIXES = ( - ".hada_w1_a", ".hada_w1_b", - ".hada_w2_a", ".hada_w2_b", - ".hada_t1", ".hada_t2", - ".alpha", ".dora_scale", ".bias", ".scale", -) -OFT_SUFFIXES = ( - ".oft_blocks", ".oft_diag", - ".alpha", ".dora_scale", ".bias", ".scale", +KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT + +BARE_DIFFUSERS_PREFIXES = ( + "layers.", "adaLN_modulation.", "final_norm.", "final_linear.", ) -LORA_MARKERS = (".lora_down.weight", ".lora_up.weight", ".lora_A.weight", ".lora_B.weight") -LOKR_MARKERS = (".lokr_w1", ".lokr_w2") -LOHA_MARKERS = (".hada_w1_a", ".hada_w1_b", ".hada_w2_a", ".hada_w2_b") -OFT_MARKERS = (".oft_blocks", ".oft_diag") -SUFFIX_NORMALIZE = { - "lora_A.weight": "lora_down.weight", - "lora_B.weight": "lora_up.weight", -} +# === Re-exports for test/back-compat === +LORA_SUFFIXES = native_loader.LORA_SUFFIXES +LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES +LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES +OFT_SUFFIXES = native_loader.OFT_SUFFIXES -def try_load_lora(name, network_on_disk, lora_scale): - """Try loading an ERNIE-Image LoRA (plus DoRA) as native modules.""" - t0 = time.time() - state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - if not has_marker(state_dict, LORA_MARKERS): - return None +LORA_MARKERS = native_loader.LORA_MARKERS +LOKR_MARKERS = native_loader.LOKR_MARKERS +LOHA_MARKERS = native_loader.LOHA_MARKERS +OFT_MARKERS = native_loader.OFT_MARKERS - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - - groups = group_by_suffixes(state_dict, LORA_SUFFIXES) - - unmapped = 0 - shape_mismatch = 0 - for network_key, w in groups.items(): - if 'lora_down.weight' not in w or 'lora_up.weight' not in w: - continue - sd_module = mapping.get(network_key) - if sd_module is None: - unmapped += 1 - continue - if not shapes_match(sd_module, w['lora_down.weight'], w['lora_up.weight']): - log.warning(f'Network load: type=LoRA name="{name}" key={network_key} shape mismatch') - shape_mismatch += 1 - continue - nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) - net.modules[network_key] = network_lora.NetworkModuleLora(net, nw) - - return finalize_network(net, name, 'LoRA', lora_scale, t0, unmapped=unmapped, mismatch=shape_mismatch) - - -def try_load_lokr(name, network_on_disk, lora_scale): - """Try loading an ERNIE-Image LoKR as native modules.""" - t0 = time.time() - state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - if not has_marker(state_dict, LOKR_MARKERS): - return None - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - - groups = group_by_suffixes(state_dict, LOKR_SUFFIXES) - - unmapped = 0 - for network_key, w in groups.items(): - has_1 = "lokr_w1" in w or ("lokr_w1_a" in w and "lokr_w1_b" in w) - has_2 = "lokr_w2" in w or ("lokr_w2_a" in w and "lokr_w2_b" in w) - if not (has_1 and has_2): - continue - sd_module = mapping.get(network_key) - if sd_module is None: - unmapped += 1 - continue - nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) - net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw) - - return finalize_network(net, name, 'LoKR', lora_scale, t0, unmapped=unmapped) - - -def try_load_loha(name, network_on_disk, lora_scale): - """Try loading an ERNIE-Image LoHA as native modules.""" - t0 = time.time() - state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - if not has_marker(state_dict, LOHA_MARKERS): - return None - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - - groups = group_by_suffixes(state_dict, LOHA_SUFFIXES) - - unmapped = 0 - for network_key, w in groups.items(): - if not all(k in w for k in ("hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b")): - continue - sd_module = mapping.get(network_key) - if sd_module is None: - unmapped += 1 - continue - nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) - net.modules[network_key] = network_hada.NetworkModuleHada(net, nw) - - return finalize_network(net, name, 'LoHA', lora_scale, t0, unmapped=unmapped) - - -def try_load_oft(name, network_on_disk, lora_scale): - """Try loading an ERNIE-Image OFT adapter as native modules.""" - t0 = time.time() - state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - if not has_marker(state_dict, OFT_MARKERS): - return None - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - - groups = group_by_suffixes(state_dict, OFT_SUFFIXES) - - unmapped = 0 - for network_key, w in groups.items(): - if not ("oft_blocks" in w or "oft_diag" in w): - continue - sd_module = mapping.get(network_key) - if sd_module is None: - unmapped += 1 - continue - nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) - net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw) - - return finalize_network(net, name, 'OFT', lora_scale, t0, unmapped=unmapped) - - -def has_marker(state_dict, markers): - return any(any(m in k for m in markers) for k in state_dict) - - -def resolve_mapping(): - sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) - lora_convert.assign_network_names_to_compvis_modules(sd_model) - return getattr(shared.sd_model, 'network_layer_mapping', {}) or {} - - -def new_network(name, network_on_disk): - net = network.Network(name, network_on_disk) - net.mtime = os.path.getmtime(network_on_disk.filename) - return net - - -def finalize_network(net, name, family, lora_scale, t0, unmapped=0, mismatch=0): - if len(net.modules) == 0: - if unmapped or mismatch: - log.debug( - f'Network load: type={family} name="{name}" native no-match' - f' unmapped={unmapped} mismatch={mismatch}' - ) - return None - log.debug( - f'Network load: type={family} name="{name}" native modules={len(net.modules)}' - f' unmapped={unmapped} mismatch={mismatch} scale={lora_scale}' - ) - l.timer.activate += time.time() - t0 - return net - - -def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool: - if not hasattr(sd_module, 'weight'): - return False - if hasattr(sd_module, 'sdnq_dequantizer'): - mod_shape = sd_module.sdnq_dequantizer.original_shape - else: - mod_shape = sd_module.weight.shape - if len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2: - return False - return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0] - - -def group_by_suffixes(state_dict, suffixes): - """Group state_dict entries by target module. - - Returns ``{network_key: {suffix: tensor, ...}}`` where ``network_key`` follows - the sdnext convention ``lora_transformer_``. Only keys - whose suffix appears in ``suffixes`` are kept; ``lora_A``/``lora_B`` are - normalized to ``lora_down``/``lora_up``. - """ - groups: dict[str, dict[str, torch.Tensor]] = {} - for key, value in state_dict.items(): - parsed = parse_key(key, suffixes) - if parsed is None: - continue - network_key, suffix = parsed - slot = groups.get(network_key) - if slot is None: - slot = {} - groups[network_key] = slot - slot[suffix] = value - return groups +SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE +BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED +has_marker = native_loader.has_marker def parse_key(key, suffixes): - stripped = key - for p in KNOWN_PREFIXES: - if key.startswith(p): - stripped = key[len(p):] - break + """ERNIE-bound :func:`native_loader.parse_key`.""" + return native_loader.parse_key( + key, suffixes, + prefixes=KNOWN_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + ) - matched_suffix = None - split_at = -1 - for marker in suffixes: - if stripped.endswith(marker): - split_at = len(stripped) - len(marker) - matched_suffix = marker.lstrip('.') - break - if split_at < 0: - return None - base = stripped[:split_at] - if not base: - return None +def group_by_suffixes(state_dict, suffixes): + """ERNIE-bound :func:`native_loader.group_by_suffixes`.""" + return native_loader.group_by_suffixes( + state_dict, suffixes, + prefixes=KNOWN_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + ) - suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix) - network_key = 'lora_transformer_' + base.replace('.', '_') - return network_key, suffix + +# === Target resolution (arch-specific) === + + +def resolve_targets(prefix_used, base): + """Passthrough for every recognized prefix. ERNIE has no fused targets or path + renames; the base path is already the diffusers module path.""" + if prefix_used in ("diffusion_model.", "transformer.", "lora_unet_", + BARE_DIFFUSERS_PREFIX_USED, None): + return [(base, None)] + return [] + + +# === Native loaders (thin wrappers over native_loader generics) === + + +_BIND_KWARGS = dict( + resolve_targets=resolve_targets, + prefixes=KNOWN_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + arch_name="ernieimage", +) + + +def try_load_lora(name, network_on_disk, lora_scale): + return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) + + +def try_load_lokr(name, network_on_disk, lora_scale): + return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) + + +def try_load_loha(name, network_on_disk, lora_scale): + return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) + + +def try_load_oft(name, network_on_disk, lora_scale): + return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) + + +def try_load(name, network_on_disk, lora_scale): + """Run every ERNIE family loader, merge any that match.""" + return native_loader.try_load_chain( + name, network_on_disk, lora_scale, + family_loaders=(try_load_lora, try_load_lokr, try_load_loha, try_load_oft), + ) diff --git a/test/test-ernie-native-adapters.py b/test/test-ernie-native-adapters.py index 8dd7fd2f5..f1332de02 100644 --- a/test/test-ernie-native-adapters.py +++ b/test/test-ernie-native-adapters.py @@ -376,21 +376,23 @@ CAT_PARSE = category('parse') def test_parse_key_all_prefixes(): - """parse_key recognizes BFL, PEFT, kohya, and bare keys.""" + """parse_key returns (prefix_used, base, suffix). ERNIE has no path renames + so resolve_targets passes the base through verbatim.""" + bd = E.BARE_DIFFUSERS_PREFIX_USED cases = [ ('diffusion_model.layers.0.mlp.gate_proj.lora_A.weight', E.LORA_SUFFIXES, - ('lora_transformer_layers_0_mlp_gate_proj', 'lora_down.weight')), + ('diffusion_model.', 'layers.0.mlp.gate_proj', 'lora_down.weight')), ('transformer.layers.1.self_attention.to_q.lora_B.weight', E.LORA_SUFFIXES, - ('lora_transformer_layers_1_self_attention_to_q', 'lora_up.weight')), + ('transformer.', 'layers.1.self_attention.to_q', 'lora_up.weight')), ('lora_unet_layers_0_mlp_linear_fc2.lora_down.weight', E.LORA_SUFFIXES, - ('lora_transformer_layers_0_mlp_linear_fc2', 'lora_down.weight')), - # Bare path (no prefix) - ernie parse_key allows fallthrough + ('lora_unet_', 'layers_0_mlp_linear_fc2', 'lora_down.weight')), + # Bare path starting with a known block prefix ('layers.0.mlp.up_proj.lora_A.weight', E.LORA_SUFFIXES, - ('lora_transformer_layers_0_mlp_up_proj', 'lora_down.weight')), + (bd, 'layers.0.mlp.up_proj', 'lora_down.weight')), ('random.unrelated.key', E.LORA_SUFFIXES, None), ] for key, suffixes, expected in cases: