mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
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.
This commit is contained in:
CalamitousFelicitousness
+87
-250
@@ -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_<path_with_underscores>``. 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),
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user