From 7c32d97121721e53198f42549d591c28ab8e5d0d Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 21:43:30 +0100 Subject: [PATCH 01/13] refactor(lora): extract native adapter scaffolding to native_loader Lifts the parts of the per-arch native loaders that are byte-identical across z-image, chroma, ernie, and flux2 into a new shared module. - Suffix and marker tables for all nine families - SUFFIX_NORMALIZE (lora_A/lora_B to lora_down/lora_up) - KNOWN_PREFIXES_DEFAULT and BARE_DIFFUSERS_PREFIX_USED sentinel - ChunkSpec dataclass for fused-weight slicing (equal idx+total or unequal start+end) - unwrap_peft_wrapper and strip_peft_adapter_name - has_marker, resolve_mapping, new_network, finalize_network, shapes_match - Parameterized parse_key and group_by_suffixes --- modules/lora/native_loader.py | 337 ++++++++++++++++++++++++++++++++++ 1 file changed, 337 insertions(+) create mode 100644 modules/lora/native_loader.py diff --git a/modules/lora/native_loader.py b/modules/lora/native_loader.py new file mode 100644 index 000000000..f399c48d7 --- /dev/null +++ b/modules/lora/native_loader.py @@ -0,0 +1,337 @@ +"""Shared scaffolding for native adapter loaders. + +The four native adapter loaders (z-image, chroma, ernie, flux2) all implement +the same algorithm: + +1. Read the safetensors state dict +2. Test for family-specific markers; bail out if absent +3. Resolve the diffusers ``network_layer_mapping`` +4. Group state-dict entries by ``(prefix, base)`` +5. For each group, ask the arch to resolve targets (diffusers paths + chunk specs) +6. Instantiate a :class:`network.NetworkModule*` per resolved target +7. Return a populated :class:`network.Network` (or ``None`` if no matches) + +This module holds the parts of that algorithm that don't vary between +architectures. Constants (suffix and marker tables), helpers (``has_marker``, +``resolve_mapping``, ``new_network``, ``finalize_network``, ``shapes_match``), +the parameterized parsing primitives (``parse_key``, ``group_by_suffixes``), +and the two cross-arch key normalizations (``unwrap_peft_wrapper`` and +``strip_peft_adapter_name``) live here. + +The variance that does remain is captured by ``ChunkSpec`` (how a row range +of a fused weight is described) and the per-arch ``resolve_targets`` callable +each loader passes in (how a parsed ``(prefix, base)`` maps to one or more +diffusers paths plus optional chunk descriptors). + +Per-arch loader modules import this module and pass their own ``prefixes``, +``bare_prefixes``, ``bare_diffusers_prefixes``, and ``resolve_targets`` to the +generic helpers. Loader business logic itself lands in subsequent commits. +""" + +import os +import time +from dataclasses import dataclass + +import torch + +from modules import shared, sd_models +from modules.logger import log +from modules.lora import lora_convert, network +from modules.lora import lora_common as l + + +# Universal prefix list shared by every native arch loader. Per-arch loaders +# extend this with arch-specific entries (e.g. flux2 adds ``"lycoris_"``). +KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_") + + +# Sentinel ``prefix_used`` value emitted by :func:`parse_key` when a bare path +# starting with a member of ``bare_diffusers_prefixes`` matches. Loader +# ``resolve_targets`` callables dispatch on this string to pass the base path +# through verbatim (no rename required, the path is already in diffusers form). +BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers" + + +SUFFIX_NORMALIZE = { + "lora_A.weight": "lora_down.weight", + "lora_B.weight": "lora_up.weight", +} + + +# === Family suffix tables === +# Alpha / scale / bias / dora_scale flow into ``weights.w`` via the base +# ``network.NetworkModule.__init__`` and are listed here so they survive the +# suffix-filter pass in :func:`parse_key`. + +LORA_SUFFIXES = ( + ".lora_down.weight", ".lora_up.weight", ".lora_mid.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", +) +IA3_SUFFIXES = ( + ".weight", ".on_input", + ".alpha", ".scale", +) +GLORA_SUFFIXES = ( + ".a1.weight", ".a2.weight", + ".b1.weight", ".b2.weight", + ".alpha", ".dora_scale", ".scale", +) +NORM_SUFFIXES = ( + ".w_norm", ".b_norm", + ".alpha", ".scale", +) +FULL_SUFFIXES = ( + ".diff", ".diff_b", + ".alpha", ".scale", +) + + +# === Family marker tables === +# Presence of any marker substring anywhere in a key triggers a try-load attempt +# for that family. Markers are deliberately narrower than suffixes (e.g. IA3's +# ``.on_input`` rather than ``.weight``) so :func:`has_marker` does not light up +# on accidental overlaps with other families. + +LORA_MARKERS = ( + ".lora_down.weight", ".lora_up.weight", + ".lora_A.weight", ".lora_B.weight", + # PEFT named-adapter saves embed the slot name as ``.lora_A..weight``; + # the trailing-dot forms catch every variant. + ".lora_A.", ".lora_B.", +) +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") +IA3_MARKERS = (".on_input",) # NOT .weight - too generic, overlaps every other family +GLORA_MARKERS = (".a1.weight", ".a2.weight", ".b1.weight", ".b2.weight") +NORM_MARKERS = (".w_norm",) +FULL_MARKERS = (".diff",) + + +# === Chunk descriptor === + + +@dataclass(frozen=True) +class ChunkSpec: + """How to slice a fused weight along dim 0 for one target module. + + Two forms supported: + + - Equal chunks (``idx`` + ``total``): fused QKV split into Q/K/V via + ``torch.chunk(up, total, dim=0)[idx]``. Used by flux2 / z-image where + Q, K and V have the same ``out_features``. + - Row range (``start`` + ``end``): asymmetric partition via + ``up[start:end]``. Used by chroma's single-block ``linear1`` which + fuses Q / K / V / proj_mlp at unequal sizes + (``[3072, 3072, 3072, 12288]``). + + Generic loaders check :attr:`is_equal_chunks` to decide between the two + forms and select the appropriate ``NetworkModule*Chunk`` / + ``NetworkModule*SliceChunk`` variant. + """ + idx: int | None = None + total: int | None = None + start: int | None = None + end: int | None = None + + @property + def is_equal_chunks(self) -> bool: + return self.idx is not None and self.total is not None + + +# === Key normalizations (applied universally by parse_key) === + + +def unwrap_peft_wrapper(key): + """Strip the ``base_model.model.`` prefix added by ``peft.save_pretrained``. + + PeftModel.save_pretrained prepends this wrapper to every adapter key. The + content underneath can be any of the standard prefixes (BFL, diffusers PEFT, + kohya, bare); a single strip lets the rest of :func:`parse_key` handle the + unwrapped key normally. Mirrors the diffusers ``Flux2LoraLoaderMixin`` + behavior of renaming ``base_model.model.`` to ``diffusion_model.`` before + feeding the converter. + """ + if key.startswith("base_model.model."): + return key[len("base_model.model."):] + return key + + +def strip_peft_adapter_name(key): + """Normalize ``.lora_[AB]..weight`` to ``.lora_[AB].weight``. + + ``peft.PeftModel`` and the diffusers ``save_lora_adapter`` exporter embed + the adapter slot name into the saved key (``"default"`` when not explicitly + set). Strip a single non-dotted name segment so the suffix table matches + without having to list every plausible adapter name. + """ + for inner in (".lora_A.", ".lora_B."): + idx = key.find(inner) + if idx == -1: + continue + rest = key[idx + len(inner):] + if rest == "weight" or not rest.endswith(".weight"): + continue + adapter_name = rest[:-len(".weight")] + if adapter_name and "." not in adapter_name: + return key[:idx] + inner + "weight" + return key + + +# === Core helpers === + + +def has_marker(state_dict, markers): + """Substring scan: does any key in ``state_dict`` contain any marker?""" + return any(any(m in k for m in markers) for k in state_dict) + + +def resolve_mapping(): + """Ensure ``network_layer_mapping`` is populated, return it (or empty dict).""" + 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): + """Construct an empty :class:`network.Network` with the file's mtime stamped.""" + 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, 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. + """ + if len(net.modules) == 0: + if unmapped or mismatch or skipped: + log.debug( + f'Network load: type={family} name="{name}" native no-match' + f' unmapped={unmapped} mismatch={mismatch} skipped={skipped}' + ) + return None + log.debug( + f'Network load: type={family} name="{name}" native modules={len(net.modules)}' + f' unmapped={unmapped} mismatch={mismatch} skipped={skipped} 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: + """LoRA-style rank-and-dim sanity check against the live module weight. + + Honors SDNQ-quantized modules by reading the original shape from the + dequantizer rather than the packed weight tensor. + """ + 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] + + +# === Parsing primitives === + + +def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=()): + """Return ``(prefix_used, base, suffix_normalized)`` or ``None``. + + ``prefix_used`` is the matched element of ``prefixes``, ``BARE_DIFFUSERS_PREFIX_USED`` + if a member of ``bare_diffusers_prefixes`` matched, or ``None`` for a key + that matched a member of ``bare_prefixes``. ``base`` is the path with prefix + and suffix removed. ``suffix_normalized`` is the suffix (without the leading + dot) after applying :data:`SUFFIX_NORMALIZE` (e.g. ``lora_A.weight`` becomes + ``lora_down.weight``). + + Always applies :func:`unwrap_peft_wrapper` and :func:`strip_peft_adapter_name` + to the raw key before format detection so callers do not have to opt in. + """ + key = unwrap_peft_wrapper(key) + key = strip_peft_adapter_name(key) + prefix_used = None + stripped = key + for p in prefixes: + if key.startswith(p): + prefix_used = p + stripped = key[len(p):] + break + if prefix_used is None: + if any(key.startswith(p) for p in bare_diffusers_prefixes): + prefix_used = BARE_DIFFUSERS_PREFIX_USED + elif not any(key.startswith(p) for p in bare_prefixes): + return None + + 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 + + suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix) + return prefix_used, base, suffix + + +def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=()): + """Group state-dict entries by ``(prefix_used, base)``. + + Returns ``{(prefix_used, base): {suffix: tensor, ...}}`` where each suffix + is the normalized form produced by :func:`parse_key`. Per-family loaders + apply their own key-presence gates on each group (e.g. LoRA requires both + ``lora_down.weight`` and ``lora_up.weight``). + """ + groups: dict[tuple, dict[str, torch.Tensor]] = {} + for key, value in state_dict.items(): + parsed = parse_key( + key, suffixes, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + if parsed is None: + continue + prefix_used, base, suffix = parsed + slot = groups.get((prefix_used, base)) + if slot is None: + slot = {} + groups[(prefix_used, base)] = slot + slot[suffix] = value + return groups + + +# Surface ``sd_models.read_state_dict`` here so loader modules don't have to +# import ``sd_models`` directly; keeps the per-arch wrapper imports compact. +read_state_dict = sd_models.read_state_dict From cbaaf1c88c2c2d063f60af0c62a496e157abcb7a Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 21:47:50 +0100 Subject: [PATCH 02/13] refactor(lora): move NetworkModuleLokrSliceChunk to network_lokr Lifts the slice variant from chroma_lora into network_lokr so the generic LoKR loader can dispatch to either NetworkModuleLokrChunk (equal chunks) or NetworkModuleLokrSliceChunk (unequal ranges) based on ChunkSpec shape. chroma_lora keeps the same slice path through an updated import. --- modules/lora/network_lokr.py | 39 ++++++++++++++++++++++++++++++ pipelines/chroma/chroma_lora.py | 42 +-------------------------------- 2 files changed, 40 insertions(+), 41 deletions(-) diff --git a/modules/lora/network_lokr.py b/modules/lora/network_lokr.py index fcb6037e3..096d7f568 100644 --- a/modules/lora/network_lokr.py +++ b/modules/lora/network_lokr.py @@ -92,3 +92,42 @@ class NetworkModuleLokrChunk(NetworkModuleLokr): updown = torch.chunk(updown, self.num_chunks, dim=0)[self.chunk_index] output_shape = list(updown.shape) return self.finalize_updown(updown, target, output_shape) + + +class NetworkModuleLokrSliceChunk(NetworkModuleLokr): + """LoKR module that returns one row-range of the Kronecker product. + + Used when a LoKR adapter targets a fused weight with unequal chunk sizes + (e.g. Chroma single ``linear1`` = Q/K/V/proj_mlp at dims + [3072, 3072, 3072, 12288]). ``NetworkModuleLokrChunk`` only supports + equal-sized chunks via ``torch.chunk``; this variant slices an explicit + row range so partitions of any shape are addressable. + """ + def __init__(self, net, weights, start_row, end_row): + super().__init__(net, weights) + self.start_row = start_row + self.end_row = end_row + + def calc_updown(self, target): + if self.w1 is not None: + w1 = self.w1.to(target.device, dtype=target.dtype) + else: + w1a = self.w1a.to(target.device, dtype=target.dtype) + w1b = self.w1b.to(target.device, dtype=target.dtype) + w1 = w1a @ w1b + if self.w2 is not None: + w2 = self.w2.to(target.device, dtype=target.dtype) + elif self.t2 is None: + w2a = self.w2a.to(target.device, dtype=target.dtype) + w2b = self.w2b.to(target.device, dtype=target.dtype) + w2 = w2a @ w2b + else: + t2 = self.t2.to(target.device, dtype=target.dtype) + w2a = self.w2a.to(target.device, dtype=target.dtype) + w2b = self.w2b.to(target.device, dtype=target.dtype) + w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b) + full_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)] + updown = make_kron(full_shape, w1, w2) + updown = updown[self.start_row:self.end_row] + output_shape = list(updown.shape) + return self.finalize_updown(updown, target, output_shape) diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 7434421a3..382bc6f3f 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -143,46 +143,6 @@ def get_block_counts(): return num_double, num_single -class NetworkModuleLokrSliceChunk(network_lokr.NetworkModuleLokr): - """LoKR module that returns one row-range of the Kronecker product. - - Used when a LoKR adapter targets a fused weight with unequal chunk sizes - (e.g., Chroma single ``linear1`` = Q/K/V/proj_mlp at dims [3072, 3072, 3072, 12288]). - The shared ``NetworkModuleLokrChunk`` only supports equal-sized chunks via - ``torch.chunk``; this variant slices an explicit row range. - """ - def __init__(self, net, weights, start_row, end_row): - super().__init__(net, weights) - self.start_row = start_row - self.end_row = end_row - - def calc_updown(self, target): - if self.w1 is not None: - w1 = self.w1.to(target.device, dtype=target.dtype) - else: - w1a = self.w1a.to(target.device, dtype=target.dtype) - w1b = self.w1b.to(target.device, dtype=target.dtype) - w1 = w1a @ w1b - if self.w2 is not None: - w2 = self.w2.to(target.device, dtype=target.dtype) - else: - from modules.lora import lyco_helpers - if self.t2 is None: - w2a = self.w2a.to(target.device, dtype=target.dtype) - w2b = self.w2b.to(target.device, dtype=target.dtype) - w2 = w2a @ w2b - else: - t2 = self.t2.to(target.device, dtype=target.dtype) - w2a = self.w2a.to(target.device, dtype=target.dtype) - w2b = self.w2b.to(target.device, dtype=target.dtype) - w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b) - full_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)] - updown = network_lokr.make_kron(full_shape, w1, w2) - updown = updown[self.start_row:self.end_row] - output_shape = list(updown.shape) - return self.finalize_updown(updown, target, output_shape) - - def try_load_lora(name, network_on_disk, lora_scale): """Try loading a Chroma LoRA (plus DoRA) as native modules.""" t0 = time.time() @@ -247,7 +207,7 @@ def try_load_lokr(name, network_on_disk, lora_scale): rng = slice_info.get(network_key) if rng is not None: start, end = rng - net.modules[network_key] = NetworkModuleLokrSliceChunk(net, nw, start, end) + net.modules[network_key] = network_lokr.NetworkModuleLokrSliceChunk(net, nw, start, end) else: net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw) From d2fd08fe3ca0226463253e81857227acbfe46995 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 21:51:29 +0100 Subject: [PATCH 03/13] refactor(lora): add generic loaders for LoRA/LoKR/LoHA/OFT Family loaders parameterized on per-arch resolve_targets callable and prefix tuples. Build network keys as "lora_transformer_" + path.replace(".", "_"). Fused-target handling: - LoRA: chunk at load time, supports both equal and unequal ChunkSpec - LoKR: dispatch to NetworkModuleLokrChunk (equal) or LokrSliceChunk (unequal), materialize kron(w1, w2) lazily - LoHA: NetworkModuleHadaChunk for equal only; Tucker-on-fused and unequal skipped with warning - OFT/BOFT: fused skipped with warning. Algorithm discriminated by oft_blocks.ndim (3-D OFT, 4-D BOFT) Plus try_load_chain umbrella for per-arch family-iteration wrappers. --- modules/lora/native_loader.py | 285 +++++++++++++++++++++++++++++++++- 1 file changed, 284 insertions(+), 1 deletion(-) diff --git a/modules/lora/native_loader.py b/modules/lora/native_loader.py index f399c48d7..cadfbdeb4 100644 --- a/modules/lora/native_loader.py +++ b/modules/lora/native_loader.py @@ -36,7 +36,10 @@ import torch from modules import shared, sd_models from modules.logger import log -from modules.lora import lora_convert, network +from modules.lora import ( + lora_convert, network, network_boft, network_hada, network_lokr, + network_lora, network_oft, +) from modules.lora import lora_common as l @@ -335,3 +338,283 @@ def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, # Surface ``sd_models.read_state_dict`` here so loader modules don't have to # import ``sd_models`` directly; keeps the per-arch wrapper imports compact. read_state_dict = sd_models.read_state_dict + + +# === Generic family loaders === +# +# Each loader takes a per-arch ``resolve_targets`` callable returning a list of +# ``(diffusers_path, ChunkSpec | None)`` tuples for each parsed group. The +# loader builds the network key as ``"lora_transformer_" + path.replace(".", "_")`` +# and instantiates the appropriate ``network.NetworkModule*`` subclass. +# +# Fused-target handling per family: +# +# - LoRA: chunk at load time (lora_up is sliced along dim 0). Both equal and +# unequal ChunkSpec shapes are supported. +# - LoKR: defer to NetworkModuleLokrChunk (equal) or NetworkModuleLokrSliceChunk +# (unequal) which materialize the Kronecker product once and return the +# designated slice. +# - LoHA: only equal ChunkSpec is supported via NetworkModuleHadaChunk. Unequal +# slices and Tucker-decomposed LoHAs on fused targets are skipped with a +# warning (no slice variant exists, and Tucker keys cannot arise on Linear +# layers per LyCORIS upstream — see network_hada.NetworkModuleHadaChunk). +# - OFT/BOFT: no chunk variant exists; fused targets are skipped with a +# warning. Discrimination is by ``oft_blocks.ndim`` (3-D OFT, 4-D BOFT), +# mirroring upstream LyCORIS ``algo_check``. + + +def _slice_lora_chunk(w, chunk: ChunkSpec): + """Return a shallow copy of ``w`` with ``lora_up.weight`` sliced per ``chunk``. + + Equal-chunks form uses ``torch.chunk`` (faster for the symmetric case); + row-range form uses tensor slicing for arbitrary partitions. + """ + up = w["lora_up.weight"] + if chunk.is_equal_chunks: + sliced = torch.chunk(up, chunk.total, dim=0)[chunk.idx].contiguous() + else: + sliced = up[chunk.start:chunk.end].contiguous() + out = dict(w) + out["lora_up.weight"] = sliced + return out + + +def try_load_lora(name, network_on_disk, lora_scale, *, + resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, + bare_prefixes=(), bare_diffusers_prefixes=(), + arch_name="generic"): + """Generic LoRA loader (handles DoRA via the universal ``finalize_updown`` hook). + + Fused targets are chunked at load time by slicing ``lora_up`` along dim 0; + the down-side is shared across the resolved targets. + """ + t0 = time.time() + state_dict = read_state_dict(network_on_disk.filename, what="network") + if not has_marker(state_dict, LORA_MARKERS): + return None + + mapping = resolve_mapping() + net = new_network(name, network_on_disk) + groups = group_by_suffixes( + state_dict, LORA_SUFFIXES, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + + unmapped = 0 + mismatch = 0 + for (prefix, base), w in groups.items(): + if "lora_down.weight" not in w or "lora_up.weight" not in w: + continue + for diffusers_path, chunk in resolve_targets(prefix, base): + network_key = "lora_transformer_" + diffusers_path.replace(".", "_") + sd_module = mapping.get(network_key) + if sd_module is None: + unmapped += 1 + continue + + target_w = _slice_lora_chunk(w, chunk) if chunk is not None else w + + if not shapes_match(sd_module, target_w["lora_down.weight"], target_w["lora_up.weight"]): + log.warning( + f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key}' + f' lora={target_w["lora_down.weight"].shape[1]}x{target_w["lora_up.weight"].shape[0]}' + f' module={getattr(sd_module, "weight", None).shape if hasattr(sd_module, "weight") else "?"}' + f' 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) + + return finalize_network(net, name, "LoRA", lora_scale, t0, unmapped=unmapped, mismatch=mismatch) + + +def try_load_lokr(name, network_on_disk, lora_scale, *, + resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, + bare_prefixes=(), bare_diffusers_prefixes=(), + arch_name="generic"): + """Generic LoKR loader. + + Stores only the compact LoKR factors and dispatches to + :class:`network_lokr.NetworkModuleLokrChunk` (equal chunks) or + :class:`network_lokr.NetworkModuleLokrSliceChunk` (row range) at apply + time. Both materialize ``kron(w1, w2)`` once per forward pass and return + the designated slice; full materialization happens lazily inside the + module rather than at load. + """ + t0 = time.time() + state_dict = 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, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + + unmapped = 0 + for (prefix, base), 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 + for diffusers_path, chunk in resolve_targets(prefix, base): + network_key = "lora_transformer_" + diffusers_path.replace(".", "_") + 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) + if chunk is None: + net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw) + elif chunk.is_equal_chunks: + net.modules[network_key] = network_lokr.NetworkModuleLokrChunk(net, nw, chunk.idx, chunk.total) + else: + net.modules[network_key] = network_lokr.NetworkModuleLokrSliceChunk(net, nw, chunk.start, chunk.end) + + return finalize_network(net, name, "LoKR", lora_scale, t0, unmapped=unmapped) + + +def try_load_loha(name, network_on_disk, lora_scale, *, + resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, + bare_prefixes=(), bare_diffusers_prefixes=(), + arch_name="generic"): + """Generic LoHA (Hadamard product) loader. + + Standard non-Tucker LoHA on fused targets uses + :class:`network_hada.NetworkModuleHadaChunk` for equal-chunks dispatch. + Tucker-decomposed LoHAs on fused targets and any unequal-chunks dispatch + are skipped with a warning: no slice variant exists, and Tucker keys + cannot arise on Linear layers per LyCORIS upstream. + """ + t0 = time.time() + state_dict = 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, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + + unmapped = 0 + skipped = 0 + for (prefix, base), 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 + is_tucker = "hada_t1" in w or "hada_t2" in w + targets = resolve_targets(prefix, base) + is_fused = any(t[1] is not None for t in targets) + if is_fused and is_tucker: + log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={base} Tucker fused QKV skipped (unsupported)') + skipped += 1 + continue + for diffusers_path, chunk in targets: + network_key = "lora_transformer_" + diffusers_path.replace(".", "_") + 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) + if chunk is None: + net.modules[network_key] = network_hada.NetworkModuleHada(net, nw) + elif chunk.is_equal_chunks: + net.modules[network_key] = network_hada.NetworkModuleHadaChunk(net, nw, chunk.idx, chunk.total) + else: + log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={network_key} unequal fused chunks unsupported') + skipped += 1 + + return finalize_network(net, name, "LoHA", lora_scale, t0, unmapped=unmapped, skipped=skipped) + + +def try_load_oft(name, network_on_disk, lora_scale, *, + resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, + bare_prefixes=(), bare_diffusers_prefixes=(), + arch_name="generic"): + """Generic OFT/BOFT loader. + + OFT and BOFT share the ``oft_blocks`` save key and are discriminated by + tensor dimensionality (3-D OFT, 4-D BOFT), mirroring upstream + ``algo_check``. Both kohya (``oft_blocks`` + alpha-as-constraint) and + LyCORIS (``oft_diag``) OFT layouts route through + :class:`network_oft.NetworkModuleOFT`; BOFT routes through + :class:`network_boft.NetworkModuleBOFT`. + + Fused targets are skipped with a warning for both algorithms: OFT block + structure (and BOFT's per-stage block partition) is tied to the target + module's ``out_features``, so a per-Q/K/V split would require re-deriving + the rotations per chunk. + """ + t0 = time.time() + state_dict = 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, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + + unmapped = 0 + skipped = 0 + for (prefix, base), w in groups.items(): + if not ("oft_blocks" in w or "oft_diag" in w): + continue + is_boft = "oft_blocks" in w and w["oft_blocks"].ndim == 4 + targets = resolve_targets(prefix, base) + if any(t[1] is not None for t in targets): + log.warning(f'Network load: type={"BOFT" if is_boft else "OFT"} name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)') + skipped += 1 + continue + for diffusers_path, _ in targets: + network_key = "lora_transformer_" + diffusers_path.replace(".", "_") + 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) + if is_boft: + net.modules[network_key] = network_boft.NetworkModuleBOFT(net, nw) + else: + net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw) + + return finalize_network(net, name, "OFT", lora_scale, t0, unmapped=unmapped, skipped=skipped) + + +# === Per-arch umbrella === + + +def try_load_chain(name, network_on_disk, lora_scale, family_loaders): + """Run each family loader in order and merge any non-None results. + + Per-arch loader modules expose a ``try_load(name, nod, scale)`` entry point + that the dispatcher in ``modules.lora.lora_load`` calls. That entry point + is a thin wrapper around this helper: it passes ``family_loaders`` as a + tuple of partial-applied generic loaders, each already bound to the arch's + ``resolve_targets`` and prefix tuples. + """ + net = None + for try_fn in family_loaders: + sub = try_fn(name, network_on_disk, lora_scale) + if sub is None: + continue + if net is None: + net = sub + else: + net.modules.update(sub.modules) + return net From 6ad67319de32a0cc2199a965b265adeedbc6cb8b Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 21:53:53 +0100 Subject: [PATCH 04/13] refactor(lora): add generic loaders for IA3/GLoRA/Norm/Full Same parameterized shape as the core 4. - IA3: .on_input is the marker disambiguator (.weight is too generic). Fused targets skipped. - GLoRA: requires a1/a2/b1/b2 per group. Fused skipped (target-dependent term doesn't slice cleanly). - Norm: never fused. Loader-local network_layer_name stamping bypasses lora_convert's transformer-norm guard without changing the carve-out. - Full: fused skipped (no chunk class for diff tensors). --- modules/lora/native_loader.py | 209 +++++++++++++++++++++++++++++++++- 1 file changed, 207 insertions(+), 2 deletions(-) diff --git a/modules/lora/native_loader.py b/modules/lora/native_loader.py index cadfbdeb4..f912a4b76 100644 --- a/modules/lora/native_loader.py +++ b/modules/lora/native_loader.py @@ -37,8 +37,9 @@ import torch from modules import shared, sd_models from modules.logger import log from modules.lora import ( - lora_convert, network, network_boft, network_hada, network_lokr, - network_lora, network_oft, + lora_convert, network, network_boft, network_full, network_glora, + network_hada, network_ia3, network_lokr, network_lora, network_norm, + network_oft, ) from modules.lora import lora_common as l @@ -596,6 +597,210 @@ def try_load_oft(name, network_on_disk, lora_scale, *, return finalize_network(net, name, "OFT", lora_scale, t0, unmapped=unmapped, skipped=skipped) +def try_load_ia3(name, network_on_disk, lora_scale, *, + resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, + bare_prefixes=(), bare_diffusers_prefixes=(), + arch_name="generic"): + """Generic IA3 loader. + + IA3 stores a per-row or per-column scale vector keyed under ``.weight`` + plus an ``.on_input`` flag selecting which axis. ``.weight`` alone is too + generic for the marker scan (it overlaps every other family's + ``.lora_down.weight`` / ``.hada_w*`` keys), so the marker gate insists on + ``.on_input`` while the suffix table includes both. + + Fused targets are skipped with a warning. There is no real-world IA3-on-DiT + prevalence to justify the asymmetry between ``on_input=True`` (which would + replicate cleanly across Q/K/V) and ``on_input=False`` (which would need + output-axis slicing). + """ + t0 = time.time() + state_dict = read_state_dict(network_on_disk.filename, what="network") + if not has_marker(state_dict, IA3_MARKERS): + return None + + mapping = resolve_mapping() + net = new_network(name, network_on_disk) + groups = group_by_suffixes( + state_dict, IA3_SUFFIXES, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + + unmapped = 0 + skipped = 0 + for (prefix, base), w in groups.items(): + if not ("weight" in w and "on_input" in w): + continue + targets = resolve_targets(prefix, base) + if any(t[1] is not None for t in targets): + log.warning(f'Network load: type=IA3 name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)') + skipped += 1 + continue + for diffusers_path, _ in targets: + network_key = "lora_transformer_" + diffusers_path.replace(".", "_") + 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_ia3.NetworkModuleIa3(net, nw) + + return finalize_network(net, name, "IA3", lora_scale, t0, unmapped=unmapped, skipped=skipped) + + +def try_load_glora(name, network_on_disk, lora_scale, *, + resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, + bare_prefixes=(), bare_diffusers_prefixes=(), + arch_name="generic"): + """Generic GLoRA loader. + + GLoRA stores four low-rank components (``a1`` / ``a2`` / ``b1`` / ``b2``) + and computes ``W_delta = w2b @ w1b + (target @ w2a) @ w1a``; the second + term is target-dependent. Fused targets are skipped: the target-dependent + term doesn't slice cleanly across projections without redirecting + calc_updown to a fused proxy weight, and the file pattern is vanishingly + rare on DiT architectures. + """ + t0 = time.time() + state_dict = read_state_dict(network_on_disk.filename, what="network") + if not has_marker(state_dict, GLORA_MARKERS): + return None + + mapping = resolve_mapping() + net = new_network(name, network_on_disk) + groups = group_by_suffixes( + state_dict, GLORA_SUFFIXES, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + + unmapped = 0 + skipped = 0 + for (prefix, base), w in groups.items(): + if not all(k in w for k in ("a1.weight", "a2.weight", "b1.weight", "b2.weight")): + continue + targets = resolve_targets(prefix, base) + if any(t[1] is not None for t in targets): + log.warning(f'Network load: type=GLoRA name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)') + skipped += 1 + continue + for diffusers_path, _ in targets: + network_key = "lora_transformer_" + diffusers_path.replace(".", "_") + 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_glora.NetworkModuleGLora(net, nw) + + return finalize_network(net, name, "GLoRA", lora_scale, t0, unmapped=unmapped, skipped=skipped) + + +def try_load_norm(name, network_on_disk, lora_scale, *, + resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, + bare_prefixes=(), bare_diffusers_prefixes=(), + arch_name="generic"): + """Generic Norm (LayerNorm / RMSNorm weight + bias delta) loader. + + Norm targets are never fused, so the chunk dispatch is dropped. + + Loader-local stamping: ``lora_convert.assign_network_names_to_compvis_modules`` + deliberately skips setting ``module.network_layer_name`` for transformer + norm modules (except SD3) because of legacy CompVis UNet collisions. This + loader bypasses the guard locally: for each target it actually binds, it + sets ``network_layer_name`` directly on the host module so + ``network_activate`` will apply the delta. Stamping is idempotent and only + touches modules a Norm adapter explicitly targets. + """ + t0 = time.time() + state_dict = read_state_dict(network_on_disk.filename, what="network") + if not has_marker(state_dict, NORM_MARKERS): + return None + + mapping = resolve_mapping() + net = new_network(name, network_on_disk) + groups = group_by_suffixes( + state_dict, NORM_SUFFIXES, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + + unmapped = 0 + for (prefix, base), w in groups.items(): + if "w_norm" not in w: + continue + targets = resolve_targets(prefix, base) + if not targets: + unmapped += 1 + continue + for diffusers_path, chunk in targets: + if chunk is not None: + continue # norm targets are not fused + network_key = "lora_transformer_" + diffusers_path.replace(".", "_") + sd_module = mapping.get(network_key) + if sd_module is None: + unmapped += 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) + + +def try_load_full(name, network_on_disk, lora_scale, *, + resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, + bare_prefixes=(), bare_diffusers_prefixes=(), + arch_name="generic"): + """Generic Full (full-rank weight delta) loader. + + Full adapters carry a complete weight delta (``diff``, same shape as the + host weight) and an optional bias delta (``diff_b``). Fused targets are + skipped with a warning: ``diff`` has the host weight's full shape and + row-slicing across three projections is well-defined arithmetically, but + no chunk class exists. + """ + t0 = time.time() + state_dict = read_state_dict(network_on_disk.filename, what="network") + if not has_marker(state_dict, FULL_MARKERS): + return None + + mapping = resolve_mapping() + net = new_network(name, network_on_disk) + groups = group_by_suffixes( + state_dict, FULL_SUFFIXES, + prefixes=prefixes, + bare_prefixes=bare_prefixes, + bare_diffusers_prefixes=bare_diffusers_prefixes, + ) + + unmapped = 0 + skipped = 0 + for (prefix, base), w in groups.items(): + if "diff" not in w: + continue + targets = resolve_targets(prefix, base) + if any(t[1] is not None for t in targets): + log.warning(f'Network load: type=Full name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)') + skipped += 1 + continue + for diffusers_path, _ in targets: + network_key = "lora_transformer_" + diffusers_path.replace(".", "_") + 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_full.NetworkModuleFull(net, nw) + + return finalize_network(net, name, "Full", lora_scale, t0, unmapped=unmapped, skipped=skipped) + + # === Per-arch umbrella === From 5ff32038f16fc70f04980eacf4fdaf224ff17d4e Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 22:02:10 +0100 Subject: [PATCH 05/13] refactor(flux2): migrate to generic native_loader Replaces flux2's eight family loaders with thin wrappers binding native_loader's generics to flux2's prefix tuples and resolve_targets. Constants and helpers (has_marker, parse_key, group_by_suffixes) are re-exported from native_loader so the offline tests address them through flux2_lora's surface unchanged. resolve_targets now returns (diffusers_path, ChunkSpec | None) tuples instead of (path, idx, num_chunks). Three parse-level tests updated. PEFT-fallback path (apply_lora_alphas, preprocess_f2_keys, apply_patch) stays in flux2_lora. --- pipelines/flux/flux2_lora.py | 833 +++++------------------------ test/test-flux2-native-adapters.py | 21 +- 2 files changed, 154 insertions(+), 700 deletions(-) diff --git a/pipelines/flux/flux2_lora.py b/pipelines/flux/flux2_lora.py index fcea16474..81994cf82 100644 --- a/pipelines/flux/flux2_lora.py +++ b/pipelines/flux/flux2_lora.py @@ -15,22 +15,13 @@ produced by ``Flux2Transformer2DModel.save_lora_adapter()``). Diffusers-PEFT ``.lora_A.default.weight``) is stripped to match the standard suffix table. BFL/kohya keys are mapped to diffusers paths via ``F2_SINGLE_MAP`` / -``F2_DOUBLE_MAP`` / ``F2_QKV_MAP``. Fused QKV in double_blocks is split into -three Q/K/V targets at lookup time. PEFT keys are diffusers paths already and -are returned verbatim with no chunking. +``F2_DOUBLE_MAP`` / ``F2_QKV_MAP``. Fused QKV in double_blocks emits three +Q/K/V targets each carrying a :class:`modules.lora.native_loader.ChunkSpec` +that the generic loaders use to chunk the up-weight or instantiate the +appropriate ``NetworkModule*Chunk`` variant. -Per-family fused-QKV handling: - -- LoRA: load-time chunk of ``lora_up`` along dim 0 (the down-side is shared). -- LoKR: apply-time slice via :class:`NetworkModuleLokrChunk`, which builds - ``kron(w1, w2)`` once and returns the designated row range. -- LoHA: apply-time slice via :class:`NetworkModuleHadaChunk`, which slices - ``w1a``/``w2a`` and computes the partial Hadamard product. Tucker - (CP-decomposed) LoHAs are not chunked and are skipped on fused targets. -- OFT, IA3, GLoRA, Full: no chunk class exists and the math is not row-sliceable - without re-deriving per-projection structure. Fused groups are skipped with a - warning. -- Norm: targets 1-D LayerNorm/RMSNorm parameters; never fused. +Per-family fused-QKV handling is inherited from +:mod:`modules.lora.native_loader`; see the loader-by-loader notes there. LyCORIS algorithm coverage relative to upstream ``KohakuBlueleaf/LyCORIS/lycoris/modules/``: @@ -56,31 +47,15 @@ to inject the ``diffusion_model.`` prefix for bare-BFL keys and bake kohya """ 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, network_boft, - network_ia3, network_glora, network_norm, network_full, lora_convert, -) -from modules.lora import lora_common as l +from modules.lora import native_loader +from modules.lora.native_loader import ChunkSpec -# === Format detection === +# === Arch-specific prefix configuration === -# Prefixes we recognize as the "true" format-identifying prefix on a state-dict -# key. The PEFT save wrapper ``base_model.model.`` is handled separately as a -# pre-strip step (see :func:`_unwrap_peft_wrapper`) because it can wrap any of -# the prefixes below — peft.save_pretrained prepends it indiscriminately. -# -# - ``diffusion_model.`` — AI-toolkit / BFL native (e.g. ostris/ai-toolkit) -# - ``transformer.`` — diffusers PEFT in-memory (e.g. HF DreamBooth scripts) -# - ``lora_unet_`` — kohya-ss/sd-scripts standard -# - ``lycoris_`` — LyCORIS-standalone save (e.g. SimpleTuner LoKR); -# the path under this prefix is an underscore-rendered -# diffusers path, not a BFL path -KNOWN_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_", "lycoris_") +KNOWN_PREFIXES = native_loader.KNOWN_PREFIXES_DEFAULT + ("lycoris_",) BARE_FLUX_PREFIXES = ( "single_blocks.", "double_blocks.", "img_in.", "txt_in.", @@ -88,749 +63,227 @@ BARE_FLUX_PREFIXES = ( "double_stream_modulation_", ) -# Bare diffusers paths (no wrapping prefix) — produced by -# ``Flux2Transformer2DModel.save_lora_adapter()`` after attaching a PEFT adapter. -# These are already-diffusers paths and pass through ``resolve_targets`` verbatim. BARE_DIFFUSERS_PREFIXES = ("single_transformer_blocks.", "transformer_blocks.") -BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers" # sentinel value for ``parse_key`` return - -SUFFIX_NORMALIZE = { - "lora_A.weight": "lora_down.weight", - "lora_B.weight": "lora_up.weight", -} -# === Family suffix tables (alpha / scale / bias / dora_scale flow into weights.w via base NetworkModule.__init__) === +# === BFL to diffusers mapping === -LORA_SUFFIXES = ( - ".lora_down.weight", ".lora_up.weight", ".lora_mid.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", -) -IA3_SUFFIXES = ( - ".weight", ".on_input", - ".alpha", ".scale", -) -GLORA_SUFFIXES = ( - ".a1.weight", ".a2.weight", - ".b1.weight", ".b2.weight", - ".alpha", ".dora_scale", ".scale", -) -NORM_SUFFIXES = ( - ".w_norm", ".b_norm", - ".alpha", ".scale", -) -FULL_SUFFIXES = ( - ".diff", ".diff_b", - ".alpha", ".scale", -) - -LORA_MARKERS = ( - ".lora_down.weight", ".lora_up.weight", - ".lora_A.weight", ".lora_B.weight", - # PEFT named-adapter saves embed the slot name as ``.lora_A..weight``; - # the trailing-dot forms catch every variant. - ".lora_A.", ".lora_B.", -) -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") -IA3_MARKERS = (".on_input",) # NOT .weight — too generic, overlaps every other family -GLORA_MARKERS = (".a1.weight", ".a2.weight", ".b1.weight", ".b2.weight") -NORM_MARKERS = (".w_norm",) -FULL_MARKERS = (".diff",) - - -# === BFL → diffusers mapping === - -# Single-block (single_transformer_blocks.{i}.) — both projections are single fused diffusers modules, -# so no chunking is needed for any adapter family. +# Single-block (single_transformer_blocks.{i}.) - both projections are +# single fused diffusers modules, so no chunking is needed for any family. F2_SINGLE_MAP = { - 'linear1': 'attn.to_qkv_mlp_proj', - 'linear2': 'attn.to_out', + "linear1": "attn.to_qkv_mlp_proj", + "linear2": "attn.to_out", } # Double-block non-QKV targets (transformer_blocks.{i}.). F2_DOUBLE_MAP = { - 'img_attn.proj': 'attn.to_out.0', - 'txt_attn.proj': 'attn.to_add_out', - 'img_mlp.0': 'ff.linear_in', - 'img_mlp.2': 'ff.linear_out', - 'txt_mlp.0': 'ff_context.linear_in', - 'txt_mlp.2': 'ff_context.linear_out', + "img_attn.proj": "attn.to_out.0", + "txt_attn.proj": "attn.to_add_out", + "img_mlp.0": "ff.linear_in", + "img_mlp.2": "ff.linear_out", + "txt_mlp.0": "ff_context.linear_in", + "txt_mlp.2": "ff_context.linear_out", } -# Double-block fused QKV targets — diffusers exposes Q/K/V as separate modules, -# so resolve_targets emits three (path, chunk_index, num_chunks=3) entries. +# Double-block fused QKV targets - diffusers exposes Q/K/V as separate modules, +# so resolve_targets emits three (path, ChunkSpec(idx, total=3)) entries. F2_QKV_MAP = { - 'img_attn.qkv': ('attn', ['to_q', 'to_k', 'to_v']), - 'txt_attn.qkv': ('attn', ['add_q_proj', 'add_k_proj', 'add_v_proj']), + "img_attn.qkv": ("attn", ["to_q", "to_k", "to_v"]), + "txt_attn.qkv": ("attn", ["add_q_proj", "add_k_proj", "add_v_proj"]), } -# Kohya underscore suffix → BFL dot suffix (last underscore becomes dot). -# Used to convert kohya key fragments to look up F2_DOUBLE_MAP / F2_QKV_MAP. +# Kohya underscore suffix -> BFL dot suffix. Used to convert kohya key fragments +# to look up F2_DOUBLE_MAP / F2_QKV_MAP. KOHYA_SUFFIX_MAP = { - 'img_attn_proj': 'img_attn.proj', - 'txt_attn_proj': 'txt_attn.proj', - 'img_attn_qkv': 'img_attn.qkv', - 'txt_attn_qkv': 'txt_attn.qkv', - 'img_mlp_0': 'img_mlp.0', - 'img_mlp_2': 'img_mlp.2', - 'txt_mlp_0': 'txt_mlp.0', - 'txt_mlp_2': 'txt_mlp.2', + "img_attn_proj": "img_attn.proj", + "txt_attn_proj": "txt_attn.proj", + "img_attn_qkv": "img_attn.qkv", + "txt_attn_qkv": "txt_attn.qkv", + "img_mlp_0": "img_mlp.0", + "img_mlp_2": "img_mlp.2", + "txt_mlp_0": "txt_mlp.0", + "txt_mlp_2": "txt_mlp.2", } -# === Shared scaffolding === +# === Re-exports for backward compatibility === +# The offline test suite addresses these via the flux2_lora module surface. +# Re-export rather than asking tests to import native_loader directly. +LORA_SUFFIXES = native_loader.LORA_SUFFIXES +LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES +LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES +OFT_SUFFIXES = native_loader.OFT_SUFFIXES +IA3_SUFFIXES = native_loader.IA3_SUFFIXES +GLORA_SUFFIXES = native_loader.GLORA_SUFFIXES +NORM_SUFFIXES = native_loader.NORM_SUFFIXES +FULL_SUFFIXES = native_loader.FULL_SUFFIXES -def has_marker(state_dict, markers): - return any(any(m in k for m in markers) for k in state_dict) +LORA_MARKERS = native_loader.LORA_MARKERS +LOKR_MARKERS = native_loader.LOKR_MARKERS +LOHA_MARKERS = native_loader.LOHA_MARKERS +OFT_MARKERS = native_loader.OFT_MARKERS +IA3_MARKERS = native_loader.IA3_MARKERS +GLORA_MARKERS = native_loader.GLORA_MARKERS +NORM_MARKERS = native_loader.NORM_MARKERS +FULL_MARKERS = native_loader.FULL_MARKERS - -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, skipped=0): - if len(net.modules) == 0: - if unmapped or mismatch or skipped: - log.debug( - f'Network load: type={family} name="{name}" native no-match' - f' unmapped={unmapped} mismatch={mismatch} skipped={skipped}' - ) - return None - log.debug( - f'Network load: type={family} name="{name}" native modules={len(net.modules)}' - f' unmapped={unmapped} mismatch={mismatch} skipped={skipped} 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 _unwrap_peft_wrapper(key): - """Strip the ``base_model.model.`` prefix added by ``peft.save_pretrained``. - - PeftModel.save_pretrained prepends this wrapper to every adapter key. The - content underneath can be any of the formats KNOWN_PREFIXES already handle: - - - BFL keys (e.g. fal/flux-2-klein-4B-outpaint-lora: - ``base_model.model.double_blocks.0.img_attn.proj.lora_A.weight``) - - Diffusers paths under ``transformer.`` (HF DreamBooth scripts that - target diffusers modules and let peft wrap them) - - Bare-BFL keys (rare but possible) - - Stripping the wrapper once is enough; the rest of :func:`parse_key` then - matches the unwrapped key against KNOWN_PREFIXES or the bare-BFL fallback - normally. Mirrors the diffusers ``Flux2LoraLoaderMixin.lora_state_dict`` - behavior at lora_pipeline.py:5684-5686, which renames the prefix to - ``diffusion_model.`` before feeding the key to the AI-toolkit converter. - """ - if key.startswith("base_model.model."): - return key[len("base_model.model."):] - return key - - -def _strip_peft_adapter_name(key): - """Normalize ``.lora_[AB]..weight`` to ``.lora_[AB].weight``. - - ``peft.PeftModel`` and the diffusers ``save_lora_adapter`` exporter embed the - adapter slot name into the saved key (``"default"`` when not explicitly - set). Strip a single non-dotted name segment so the suffix table matches - without listing every plausible adapter name. - """ - for inner in (".lora_A.", ".lora_B."): - idx = key.find(inner) - if idx == -1: - continue - rest = key[idx + len(inner):] - if rest == "weight" or not rest.endswith(".weight"): - continue - adapter_name = rest[:-len(".weight")] - if adapter_name and "." not in adapter_name: - return key[:idx] + inner + "weight" - return key +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): - """Return ``(prefix_used, base, suffix_normalized)`` or ``None``. - - ``prefix_used`` is the matched ``KNOWN_PREFIXES`` element, or ``None`` for - bare BFL keys. ``base`` is the format-native module path (kohya / lycoris - underscore-style or BFL / diffusers dot-style depending on prefix). - """ - key = _unwrap_peft_wrapper(key) - key = _strip_peft_adapter_name(key) - prefix_used = None - stripped = key - for p in KNOWN_PREFIXES: - if key.startswith(p): - prefix_used = p - stripped = key[len(p):] - break - if prefix_used is None: - if any(key.startswith(p) for p in BARE_DIFFUSERS_PREFIXES): - prefix_used = BARE_DIFFUSERS_PREFIX_USED - elif not any(key.startswith(p) for p in BARE_FLUX_PREFIXES): - return None - - 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 - - suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix) - return prefix_used, base, suffix + """Flux2-bound :func:`native_loader.parse_key`. Returns ``(prefix_used, base, suffix)`` or ``None``.""" + return native_loader.parse_key( + key, suffixes, + prefixes=KNOWN_PREFIXES, + bare_prefixes=BARE_FLUX_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + ) def group_by_suffixes(state_dict, suffixes): - """Group state_dict entries by ``(prefix_used, base)``. + """Flux2-bound :func:`native_loader.group_by_suffixes`.""" + return native_loader.group_by_suffixes( + state_dict, suffixes, + prefixes=KNOWN_PREFIXES, + bare_prefixes=BARE_FLUX_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + ) - Returns ``{(prefix_used, base): {suffix: tensor, ...}}`` where - ``prefix_used`` is a ``KNOWN_PREFIXES`` element or ``None`` for bare-BFL. - Per-family loaders apply their own key-presence gates on each group. - """ - groups: dict[tuple, dict[str, torch.Tensor]] = {} - for key, value in state_dict.items(): - parsed = parse_key(key, suffixes) - if parsed is None: - continue - prefix_used, base, suffix = parsed - slot = groups.get((prefix_used, base)) - if slot is None: - slot = {} - groups[(prefix_used, base)] = slot - slot[suffix] = value - return groups + +# === Target resolution (arch-specific) === def resolve_targets(prefix_used, base): - """Return ``[(diffusers_path, chunk_index, num_chunks), ...]`` for a parsed group key. + """Return ``[(diffusers_path, ChunkSpec | None), ...]`` for a parsed group key. - For kohya prefix, applies ``KOHYA_SUFFIX_MAP`` then ``F2_*_MAP``. For - BFL/bare-BFL, applies ``F2_*_MAP`` directly. For PEFT (``transformer.``), - returns the base verbatim with no chunking — it is already a diffusers path. + For ``lora_unet_`` prefix, applies ``KOHYA_SUFFIX_MAP`` then ``F2_*_MAP``. + For BFL / bare-BFL, applies ``F2_*_MAP`` directly. For ``transformer.``, + ``lycoris_``, and bare-diffusers, returns the base verbatim with no chunking. + Unrecognized prefixes return an empty list. """ - if prefix_used == 'lora_unet_': + if prefix_used == "lora_unet_": return _kohya_to_diffusers_targets(base) - if prefix_used in (None, 'diffusion_model.'): + if prefix_used in (None, "diffusion_model."): return _bfl_to_diffusers_targets(base) - if prefix_used == 'transformer.': - return [(base, None, None)] + if prefix_used == "transformer.": + return [(base, None)] if prefix_used == BARE_DIFFUSERS_PREFIX_USED: - # Already-diffusers path with no wrapping prefix (e.g. produced by - # Flux2Transformer2DModel.save_lora_adapter()). Pass through verbatim. - return [(base, None, None)] - if prefix_used == 'lycoris_': + return [(base, None)] + if prefix_used == "lycoris_": # base is an already-underscored diffusers path (e.g. # 'transformer_blocks_0_attn_add_k_proj'). The caller's network_key # construction does base.replace('.', '_'); for already-underscored # paths that's a no-op, so the network_key matches the entry stamped # by lora_convert.assign_network_names_to_compvis_modules # (e.g. 'lora_transformer_transformer_blocks_0_attn_add_k_proj'). - return [(base, None, None)] + return [(base, None)] return [] def _kohya_to_diffusers_targets(stripped): """For kohya keys like ``double_blocks_0_img_attn_proj`` or ``single_blocks_5_linear1``.""" - targets: list[tuple[str, int | None, int | None]] = [] - if stripped.startswith('single_blocks_'): - rest = stripped[len('single_blocks_'):] - idx, _, suffix = rest.partition('_') + targets: list[tuple[str, ChunkSpec | None]] = [] + if stripped.startswith("single_blocks_"): + rest = stripped[len("single_blocks_"):] + idx, _, suffix = rest.partition("_") if suffix in F2_SINGLE_MAP: - targets.append((f'single_transformer_blocks.{idx}.{F2_SINGLE_MAP[suffix]}', None, None)) - elif stripped.startswith('double_blocks_'): - rest = stripped[len('double_blocks_'):] - idx, _, kohya_suffix = rest.partition('_') + targets.append((f"single_transformer_blocks.{idx}.{F2_SINGLE_MAP[suffix]}", None)) + elif stripped.startswith("double_blocks_"): + rest = stripped[len("double_blocks_"):] + idx, _, kohya_suffix = rest.partition("_") bfl_suffix = KOHYA_SUFFIX_MAP.get(kohya_suffix) if bfl_suffix is None: return targets if bfl_suffix in F2_DOUBLE_MAP: - targets.append((f'transformer_blocks.{idx}.{F2_DOUBLE_MAP[bfl_suffix]}', None, None)) + targets.append((f"transformer_blocks.{idx}.{F2_DOUBLE_MAP[bfl_suffix]}", None)) elif bfl_suffix in F2_QKV_MAP: attn_prefix, proj_keys = F2_QKV_MAP[bfl_suffix] for i, proj_key in enumerate(proj_keys): - targets.append((f'transformer_blocks.{idx}.{attn_prefix}.{proj_key}', i, len(proj_keys))) + targets.append(( + f"transformer_blocks.{idx}.{attn_prefix}.{proj_key}", + ChunkSpec(idx=i, total=len(proj_keys)), + )) return targets def _bfl_to_diffusers_targets(base): """For BFL keys like ``double_blocks.0.img_attn.proj`` or ``single_blocks.5.linear1``.""" - targets: list[tuple[str, int | None, int | None]] = [] - parts = base.split('.') + targets: list[tuple[str, ChunkSpec | None]] = [] + parts = base.split(".") if len(parts) < 3: return targets - block_type, block_idx, module_suffix = parts[0], parts[1], '.'.join(parts[2:]) - if block_type == 'single_blocks' and module_suffix in F2_SINGLE_MAP: - targets.append((f'single_transformer_blocks.{block_idx}.{F2_SINGLE_MAP[module_suffix]}', None, None)) - elif block_type == 'double_blocks': + block_type, block_idx, module_suffix = parts[0], parts[1], ".".join(parts[2:]) + if block_type == "single_blocks" and module_suffix in F2_SINGLE_MAP: + targets.append((f"single_transformer_blocks.{block_idx}.{F2_SINGLE_MAP[module_suffix]}", None)) + elif block_type == "double_blocks": if module_suffix in F2_DOUBLE_MAP: - targets.append((f'transformer_blocks.{block_idx}.{F2_DOUBLE_MAP[module_suffix]}', None, None)) + targets.append((f"transformer_blocks.{block_idx}.{F2_DOUBLE_MAP[module_suffix]}", None)) elif module_suffix in F2_QKV_MAP: attn_prefix, proj_keys = F2_QKV_MAP[module_suffix] for i, proj_key in enumerate(proj_keys): - targets.append((f'transformer_blocks.{block_idx}.{attn_prefix}.{proj_key}', i, len(proj_keys))) + targets.append(( + f"transformer_blocks.{block_idx}.{attn_prefix}.{proj_key}", + ChunkSpec(idx=i, total=len(proj_keys)), + )) return targets -# === Native loaders === +# === Native loaders (thin wrappers over native_loader generics) === -def try_load(name, network_on_disk, lora_scale): - """Run every Flux2 family loader in dispatch order, merge any that match. - - Per-family ``try_load_*`` entry points stay public; this is the single - umbrella the dispatcher in ``modules.lora.lora_load.load_safetensors`` - calls. Order matters only for marker-cost: LoRA / LoKR are most common - so their fast bail-out runs first; the rare families come last. - - Returns a ``Network`` with the union of modules from every matching - family loader, or ``None`` if no loader recognized the file. - """ - net = None - for try_fn in ( - try_load_lora, try_load_lokr, try_load_loha, try_load_oft, - try_load_ia3, try_load_glora, try_load_norm, try_load_full, - ): - sub = try_fn(name, network_on_disk, lora_scale) - if sub is None: - continue - if net is None: - net = sub - else: - net.modules.update(sub.modules) - return net +_BIND_KWARGS = dict( + resolve_targets=resolve_targets, + prefixes=KNOWN_PREFIXES, + bare_prefixes=BARE_FLUX_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + arch_name="f2", +) def try_load_lora(name, network_on_disk, lora_scale): - """Load a Flux2/Klein LoRA (plus DoRA via the universal ``finalize_updown`` hook) as native modules. - - Handles kohya, AI-toolkit/BFL, diffusers PEFT, and bare-BFL key formats. - Fused QKV in double_blocks is split at load time by chunking the up-weight - along dim 0; the down-weight is shared across Q/K/V. - """ - 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 - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - groups = group_by_suffixes(state_dict, LORA_SUFFIXES) - - unmapped = 0 - mismatch = 0 - for (prefix, base), w in groups.items(): - if 'lora_down.weight' not in w or 'lora_up.weight' not in w: - continue - for diffusers_path, chunk_idx, num_chunks in resolve_targets(prefix, base): - network_key = "lora_transformer_" + diffusers_path.replace(".", "_") - sd_module = mapping.get(network_key) - if sd_module is None: - unmapped += 1 - continue - - if chunk_idx is not None: - chunks = torch.chunk(w['lora_up.weight'], num_chunks, dim=0) - target_w = dict(w) - target_w['lora_up.weight'] = chunks[chunk_idx].contiguous() - else: - target_w = w - - if not shapes_match(sd_module, target_w['lora_down.weight'], target_w['lora_up.weight']): - log.warning( - f'Network load: type=LoRA name="{name}" key={network_key}' - f' lora={target_w["lora_down.weight"].shape[1]}x{target_w["lora_up.weight"].shape[0]}' - f' module={getattr(sd_module, "weight", None).shape if hasattr(sd_module, "weight") else "?"}' - f' 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) - - return finalize_network(net, name, 'LoRA', lora_scale, t0, unmapped=unmapped, mismatch=mismatch) + return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): - """Load a Flux2/Klein LoKR as native modules. - - Stores only the compact LoKR factors (``w1``/``w2``) and computes - ``kron(w1, w2)`` on-the-fly during weight application. For fused QKV - targets in double_blocks, :class:`NetworkModuleLokrChunk` materializes the - full Kronecker product and returns the designated Q/K/V slice. - """ - 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 (prefix, base), 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 - for diffusers_path, chunk_idx, num_chunks in resolve_targets(prefix, base): - network_key = "lora_transformer_" + diffusers_path.replace(".", "_") - 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) - if chunk_idx is not None: - net.modules[network_key] = network_lokr.NetworkModuleLokrChunk(net, nw, chunk_idx, num_chunks) - else: - net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw) - - return finalize_network(net, name, 'LoKR', lora_scale, t0, unmapped=unmapped) + return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_loha(name, network_on_disk, lora_scale): - """Load a Flux2/Klein LoHA (Hadamard product) adapter as native modules. - - Standard non-Tucker LoHA on fused QKV in double_blocks is supported via - :class:`NetworkModuleHadaChunk`, which slices ``w1a``/``w2a`` at the - chunk's row range and computes the partial Hadamard. Tucker - (CP-decomposed) LoHAs are skipped on fused targets because the chunk - class does not implement the CP path; non-fused Tucker LoHAs go through - the standard :class:`NetworkModuleHada`. - """ - 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 - skipped = 0 - for (prefix, base), 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 - is_tucker = 'hada_t1' in w or 'hada_t2' in w - targets = resolve_targets(prefix, base) - is_fused = any(t[1] is not None for t in targets) - if is_fused and is_tucker: - log.warning(f'Network load: type=LoHA name="{name}" key={base} Tucker fused QKV skipped (unsupported)') - skipped += 1 - continue - for diffusers_path, chunk_idx, num_chunks in targets: - network_key = "lora_transformer_" + diffusers_path.replace(".", "_") - 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) - if chunk_idx is not None: - net.modules[network_key] = network_hada.NetworkModuleHadaChunk(net, nw, chunk_idx, num_chunks) - else: - net.modules[network_key] = network_hada.NetworkModuleHada(net, nw) - - return finalize_network(net, name, 'LoHA', lora_scale, t0, unmapped=unmapped, skipped=skipped) + return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_oft(name, network_on_disk, lora_scale): - """Load a Flux2/Klein OFT or BOFT adapter as native modules. - - Both algorithms share the ``oft_blocks`` save key and are discriminated - by tensor dimensionality, mirroring LyCORIS's own ``algo_check``: - - - **OFT** — 3-D ``(num_blocks, block_size, block_size)``. Both kohya - (``oft_blocks`` + alpha-as-constraint) and LyCORIS (``oft_diag``) - layouts route through :class:`NetworkModuleOFT`. - - **BOFT** — 4-D ``(boft_m, block_num, block_size, block_size)``, - a cascade of butterfly factors. Routes through - :class:`NetworkModuleBOFT` which ports the butterfly-cascade - ``make_weight`` from LyCORIS boft.py. - - Fused QKV in double_blocks is skipped with a warning for both: an OFT - block structure (and BOFT's per-stage block partition) is tied to the - target module's ``out_features``, so a per-Q/K/V split would require - re-deriving the rotations per chunk. Single-block ``linear1`` (a single - fused diffusers module) and all non-QKV double-block targets work fully. - """ - 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 - skipped = 0 - for (prefix, base), w in groups.items(): - if not ('oft_blocks' in w or 'oft_diag' in w): - continue - is_boft = 'oft_blocks' in w and w['oft_blocks'].ndim == 4 - targets = resolve_targets(prefix, base) - if any(t[1] is not None for t in targets): - log.warning(f'Network load: type={"BOFT" if is_boft else "OFT"} name="{name}" key={base} fused QKV skipped (unsupported)') - skipped += 1 - continue - for diffusers_path, _, _ in targets: - network_key = "lora_transformer_" + diffusers_path.replace(".", "_") - 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) - if is_boft: - net.modules[network_key] = network_boft.NetworkModuleBOFT(net, nw) - else: - net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw) - - return finalize_network(net, name, 'OFT', lora_scale, t0, unmapped=unmapped, skipped=skipped) + return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_ia3(name, network_on_disk, lora_scale): - """Load a Flux2/Klein IA3 adapter as native modules. - - IA3 stores a per-row or per-column scale vector keyed under ``.weight`` - plus an ``.on_input`` flag selecting which axis. The ``.on_input`` marker - is the format disambiguator — ``.weight`` alone is too generic and - overlaps every other family's ``.lora_down.weight`` / ``.hada_w*`` keys, - so the SUFFIXES table includes it but the MARKERS gate insists on - ``.on_input``. - - Fused QKV in double_blocks is skipped: ``on_input=True`` IA3 vectors - would replicate cleanly to Q/K/V (same ``in_features``) but - ``on_input=False`` requires slicing the output-axis vector across the - three projections, and there is zero real-world IA3-on-DiT prevalence to - justify the asymmetry. - """ - t0 = time.time() - state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - if not has_marker(state_dict, IA3_MARKERS): - return None - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - groups = group_by_suffixes(state_dict, IA3_SUFFIXES) - - unmapped = 0 - skipped = 0 - for (prefix, base), w in groups.items(): - if not ('weight' in w and 'on_input' in w): - continue - targets = resolve_targets(prefix, base) - if any(t[1] is not None for t in targets): - log.warning(f'Network load: type=IA3 name="{name}" key={base} fused QKV skipped (unsupported)') - skipped += 1 - continue - for diffusers_path, _, _ in targets: - network_key = "lora_transformer_" + diffusers_path.replace(".", "_") - 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_ia3.NetworkModuleIa3(net, nw) - - return finalize_network(net, name, 'IA3', lora_scale, t0, unmapped=unmapped, skipped=skipped) + return native_loader.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_glora(name, network_on_disk, lora_scale): - """Load a Flux2/Klein GLoRA adapter as native modules. - - GLoRA stores four low-rank components (``a1``/``a2``/``b1``/``b2``) and - computes ``ΔW = w2b @ w1b + (target @ w2a) @ w1a`` — the second term is - target-dependent. Fused QKV in double_blocks is skipped with a warning - because the target-dependent term doesn't slice cleanly without - redirecting calc_updown to a fused proxy weight, and zero real-world - GLoRA-on-DiT files exist. - - Depends on the ``self.dim`` initialization fix in network_glora.py so - that alpha-based ``calc_scale`` is honored. - """ - t0 = time.time() - state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - if not has_marker(state_dict, GLORA_MARKERS): - return None - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - groups = group_by_suffixes(state_dict, GLORA_SUFFIXES) - - unmapped = 0 - skipped = 0 - for (prefix, base), w in groups.items(): - if not all(k in w for k in ('a1.weight', 'a2.weight', 'b1.weight', 'b2.weight')): - continue - targets = resolve_targets(prefix, base) - if any(t[1] is not None for t in targets): - log.warning(f'Network load: type=GLoRA name="{name}" key={base} fused QKV skipped (unsupported)') - skipped += 1 - continue - for diffusers_path, _, _ in targets: - network_key = "lora_transformer_" + diffusers_path.replace(".", "_") - 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_glora.NetworkModuleGLora(net, nw) - - return finalize_network(net, name, 'GLoRA', lora_scale, t0, unmapped=unmapped, skipped=skipped) + return native_loader.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_norm(name, network_on_disk, lora_scale): - """Load a Flux2/Klein Norm adapter (LayerNorm/RMSNorm weight + bias deltas) as native modules. - - Norm adapters target the RMSNorm modules inside Flux2 attention - (``attn.norm_q``, ``attn.norm_k``, ``attn.norm_added_q``, - ``attn.norm_added_k``) — the only norm modules in Flux2 with trainable - weights. The block-level ``norm1``/``norm2`` LayerNorms have - ``elementwise_affine=False`` and are not adaptable. - - Loader-local stamping: ``modules/lora/lora_convert.py:assign_network_names_to_compvis_modules`` - deliberately skips setting ``module.network_layer_name`` for transformer - norm modules (except SD3) because of legacy CompVis UNet collisions. This - loader bypasses the guard locally — for each target it actually binds, it - sets ``network_layer_name`` directly on the host module so - ``network_activate`` will apply the delta. No edit to the shared - ``lora_convert`` carve-out is required, and no norm module is touched - unless a Norm adapter explicitly targets it. - - BFL/kohya prefix support is deferred — there is no public Flux2 BFL norm - mapping table to verify against. PEFT prefix (the format produced by - ``peft`` training) works directly because the base path is already a - diffusers path. - """ - t0 = time.time() - state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - if not has_marker(state_dict, NORM_MARKERS): - return None - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - groups = group_by_suffixes(state_dict, NORM_SUFFIXES) - - unmapped = 0 - for (prefix, base), w in groups.items(): - if 'w_norm' not in w: - continue - targets = resolve_targets(prefix, base) - if not targets: - unmapped += 1 - continue - for diffusers_path, chunk_idx, _ in targets: - if chunk_idx is not None: - continue # norm targets are not fused - network_key = "lora_transformer_" + diffusers_path.replace(".", "_") - sd_module = mapping.get(network_key) - if sd_module is None: - unmapped += 1 - continue - # Bypass the lora_convert.py:502 transformer-norm guard locally. - # Stamping is idempotent and only touches modules a Norm adapter targets. - 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 native_loader.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS) def try_load_full(name, network_on_disk, lora_scale): - """Load a Flux2/Klein Full (full-rank) adapter as native modules. + return native_loader.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS) - Full adapters carry a complete weight delta (``diff``, same shape as the - host weight) and an optional bias delta (``diff_b``) via - :class:`NetworkModuleFull`. Most realistic use: small per-block bias-only - adjustments in distillation LoRAs. - Fused QKV in double_blocks is skipped with a warning. Full's ``diff`` has - the host weight's full shape; row-slicing across three projections is - well-defined arithmetically but no chunk class exists and zero - real-world Full-on-fused-DiT files exist. Single-block linear1 (a single - fused diffusers module) and non-QKV double-block targets work fully. - """ - t0 = time.time() - state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') - if not has_marker(state_dict, FULL_MARKERS): - return None - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - groups = group_by_suffixes(state_dict, FULL_SUFFIXES) - - unmapped = 0 - skipped = 0 - for (prefix, base), w in groups.items(): - if 'diff' not in w: - continue - targets = resolve_targets(prefix, base) - if any(t[1] is not None for t in targets): - log.warning(f'Network load: type=Full name="{name}" key={base} fused QKV skipped (unsupported)') - skipped += 1 - continue - for diffusers_path, _, _ in targets: - network_key = "lora_transformer_" + diffusers_path.replace(".", "_") - 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_full.NetworkModuleFull(net, nw) - - return finalize_network(net, name, 'Full', lora_scale, t0, unmapped=unmapped, skipped=skipped) +def try_load(name, network_on_disk, lora_scale): + """Single dispatcher entry point: run every 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, + try_load_ia3, try_load_glora, try_load_norm, try_load_full, + ), + ) # === Diffusers-PEFT path helpers (used when lora_force_diffusers is on) === @@ -845,12 +298,12 @@ def apply_lora_alphas(state_dict): causing a ``ValueError`` on leftover keys. This matches the approach used by ``_convert_kohya_flux_lora_to_diffusers`` for Flux 1. """ - alpha_keys = [k for k in state_dict if k.endswith('.alpha')] + alpha_keys = [k for k in state_dict if k.endswith(".alpha")] if not alpha_keys: return state_dict for alpha_key in alpha_keys: - base = alpha_key[:-len('.alpha')] - down_key = f'{base}.lora_down.weight' + base = alpha_key[:-len(".alpha")] + down_key = f"{base}.lora_down.weight" if down_key not in state_dict: continue down_weight = state_dict[down_key] @@ -863,10 +316,10 @@ def apply_lora_alphas(state_dict): scale_down *= 2 scale_up /= 2 state_dict[down_key] = down_weight * scale_down - up_key = f'{base}.lora_up.weight' + up_key = f"{base}.lora_up.weight" if up_key in state_dict: state_dict[up_key] = state_dict[up_key] * scale_up - remaining = [k for k in state_dict if k.endswith('.alpha')] + remaining = [k for k in state_dict if k.endswith(".alpha")] if remaining: log.debug(f'Network load: type=LoRA stripped {len(remaining)} orphaned alpha keys') for k in remaining: @@ -880,7 +333,7 @@ def preprocess_f2_keys(state_dict): if any(k.startswith("diffusion_model.") or k.startswith("base_model.model.") for k in state_dict): return state_dict if any(k.startswith(p) for k in state_dict for p in BARE_FLUX_PREFIXES): - log.debug('Network load: type=LoRA adding diffusion_model prefix for bare BFL-format keys') + log.debug("Network load: type=LoRA adding diffusion_model prefix for bare BFL-format keys") state_dict = {f"diffusion_model.{k}": v for k, v in state_dict.items()} return state_dict @@ -911,13 +364,13 @@ def apply_patch(): pretrained_model_name_or_path_or_dict = apply_lora_alphas(pretrained_model_name_or_path_or_dict) elif isinstance(pretrained_model_name_or_path_or_dict, (str, os.PathLike)): path = str(pretrained_model_name_or_path_or_dict) - if path.endswith('.safetensors'): + if path.endswith(".safetensors"): try: from safetensors import safe_open with safe_open(path, framework="pt") as f: keys = list(f.keys()) needs_load = ( - any(k.endswith('.alpha') for k in keys) + any(k.endswith(".alpha") for k in keys) or (not any(k.startswith("diffusion_model.") or k.startswith("base_model.model.") for k in keys) and any(k.startswith(p) for k in keys for p in BARE_FLUX_PREFIXES)) ) diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py index c84546943..56a8fa9fd 100644 --- a/test/test-flux2-native-adapters.py +++ b/test/test-flux2-native-adapters.py @@ -580,26 +580,27 @@ def test_parse_key_all_prefixes(): def test_resolve_targets_qkv_chunking(): + from modules.lora.native_loader import ChunkSpec # Kohya double_blocks fused QKV → three chunks targeting Q/K/V. targets = F.resolve_targets('lora_unet_', 'double_blocks_0_img_attn_qkv') assert targets == [ - ('transformer_blocks.0.attn.to_q', 0, 3), - ('transformer_blocks.0.attn.to_k', 1, 3), - ('transformer_blocks.0.attn.to_v', 2, 3), + ('transformer_blocks.0.attn.to_q', ChunkSpec(idx=0, total=3)), + ('transformer_blocks.0.attn.to_k', ChunkSpec(idx=1, total=3)), + ('transformer_blocks.0.attn.to_v', ChunkSpec(idx=2, total=3)), ], f'kohya img_attn.qkv → {targets}' targets = F.resolve_targets('lora_unet_', 'double_blocks_5_txt_attn_qkv') assert targets == [ - ('transformer_blocks.5.attn.add_q_proj', 0, 3), - ('transformer_blocks.5.attn.add_k_proj', 1, 3), - ('transformer_blocks.5.attn.add_v_proj', 2, 3), + ('transformer_blocks.5.attn.add_q_proj', ChunkSpec(idx=0, total=3)), + ('transformer_blocks.5.attn.add_k_proj', ChunkSpec(idx=1, total=3)), + ('transformer_blocks.5.attn.add_v_proj', ChunkSpec(idx=2, total=3)), ], f'kohya txt_attn.qkv → {targets}' targets = F.resolve_targets('diffusion_model.', 'single_blocks.7.linear1') - assert targets == [('single_transformer_blocks.7.attn.to_qkv_mlp_proj', None, None)] + assert targets == [('single_transformer_blocks.7.attn.to_qkv_mlp_proj', None)] targets = F.resolve_targets('transformer.', 'transformer_blocks.0.attn.to_q') - assert targets == [('transformer_blocks.0.attn.to_q', None, None)] + assert targets == [('transformer_blocks.0.attn.to_q', None)] targets = F.resolve_targets('weird_prefix.', 'whatever') assert targets == [] @@ -654,7 +655,7 @@ def test_parse_key_lycoris_prefix(): # resolve_targets: the underscored path is returned verbatim (no chunk). targets = F.resolve_targets('lycoris_', 'transformer_blocks_0_attn_add_k_proj') - assert targets == [('transformer_blocks_0_attn_add_k_proj', None, None)], f'targets={targets}' + assert targets == [('transformer_blocks_0_attn_add_k_proj', None)], f'targets={targets}' return True @@ -691,7 +692,7 @@ def test_parse_key_bare_diffusers_and_peft_default(): # resolve_targets passes the bare-diffusers path through verbatim. targets = F.resolve_targets(bd, 'single_transformer_blocks.5.attn.to_out') - assert targets == [('single_transformer_blocks.5.attn.to_out', None, None)], f'targets={targets}' + assert targets == [('single_transformer_blocks.5.attn.to_out', None)], f'targets={targets}' return True From 409a30f9e0d12659adab31208dd76e4b9c642b4f Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 22:21:47 +0100 Subject: [PATCH 06/13] test(zimage): offline tests for native adapter loaders Covers z-image's four-family surface (LoRA, LoKR, LoHA, OFT). Mock transformer matches diffusers.ZImageTransformer2DModel (layers, noise_refiner, context_refiner). Formats exercised: - BFL / AI-toolkit - PEFT - kohya - Legacy fused attention.qkv (split by expand_legacy_attention_*) Covers parse primitives, all four loader entry points, DoRA threading, calc_updown shape sanity per family. --- test/test-zimage-native-adapters.py | 723 ++++++++++++++++++++++++++++ 1 file changed, 723 insertions(+) create mode 100644 test/test-zimage-native-adapters.py diff --git a/test/test-zimage-native-adapters.py b/test/test-zimage-native-adapters.py new file mode 100644 index 000000000..b75a0ab00 --- /dev/null +++ b/test/test-zimage-native-adapters.py @@ -0,0 +1,723 @@ +#!/usr/bin/env python +""" +Offline unit tests for Z-Image native adapter loaders. + +Covers the four native families currently supported by ``pipelines.z_image.zimage_lora`` +(LoRA, LoKR, LoHA, OFT) plus DoRA threading via the universal +``NetworkModule.finalize_updown`` hook. + +Tests build a mock Z-Image-shaped transformer, write synthetic safetensors +files for each adapter format observed in the wild, and exercise the full +loader path from state dict to ``NetworkModule*`` instantiation. + +Save formats are cross-referenced against real Z-Image LoRAs in the wild: + +- BFL / AI-toolkit (``diffusion_model.layers.0.attention.to_q.lora_A.weight``): + e.g. ``80sFantasyZBase``, ``zimagebase_blending_v1`` +- kohya (``lora_unet_layers_0_attention_to_q.lora_down.weight``): pattern + produced by kohya-ss/sd-scripts targeting Z-Image +- PEFT (``transformer.layers.0.attention.to_q.lora_down.weight``): + e.g. ``FameGrid_Revolution_ZIB_BOLD`` +- Legacy fused ``attention.qkv`` (Z-Image pre-refactor): the loader splits + the up-weight along dim 0 at load time and emits Q / K / V targets + +Loader correctness is verified against the documented LyCORIS / kohya +on-disk shapes. The Z-Image diffusers transformer layout +(``layers[i].attention.{to_q,to_k,to_v,to_out[0]}`` plus +``feed_forward.net.{0.proj,2}`` and ``adaLN_modulation.0``) is taken +straight from ``diffusers.ZImageTransformer2DModel`` / +``ZImageTransformerBlock``. + +No running server required. + +Usage: + python test/test-zimage-native-adapters.py +""" + +import os +import sys +import tempfile +import time + +import torch +import safetensors.torch + +script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, script_dir) +os.chdir(script_dir) + +os.environ['SD_INSTALL_QUIET'] = '1' + +# Bootstrap cmd_args before any module that pulls in shared.py. +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +_orig_argv = sys.argv +sys.argv = [sys.argv[0]] +try: + modules.cmd_args.parse_args() +finally: + sys.argv = _orig_argv +installer.add_args(modules.cmd_args.parser) +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 # pylint: disable=wrong-import-position +from modules.lora import ( # pylint: disable=wrong-import-position + network, network_lora, network_lokr, network_hada, network_oft, +) +from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position +from pipelines.z_image import zimage_lora as Z # pylint: disable=wrong-import-position + + +# ============================================================ +# Test infrastructure +# ============================================================ + +results: dict[str, dict] = {} + + +def category(name: str): + if name not in results: + results[name] = {'passed': 0, 'failed': 0, 'tests': []} + return name + + +def record(cat: str, passed: bool, name: str, detail: str = ''): + status = 'PASS' if passed else 'FAIL' + results[cat]['passed' if passed else 'failed'] += 1 + results[cat]['tests'].append((status, name)) + msg = f' {status}: {name}' + if detail: + msg += f' ({detail})' + if passed: + log.info(msg) + else: + log.error(msg) + + +def run_test(cat: str, fn): + name = fn.__name__ + try: + ok = fn() + if ok is False: + record(cat, False, name) + else: + record(cat, True, name) + except AssertionError as e: + record(cat, False, name, str(e)) + except Exception as e: # pylint: disable=broad-except + record(cat, False, name, f'exception: {e}') + import traceback + traceback.print_exc() + + +# ============================================================ +# Mock Z-Image transformer +# ============================================================ +# Shape constants chosen to mirror real Z-Image proportions while keeping +# tensors small: ``hidden_dim = int(dim / 3 * 8)`` matches the +# ZImageTransformerBlock FeedForward sizing, and adaLN_modulation outputs +# ``4 * dim`` per upstream. + +HIDDEN = 96 # dim +HEAD_DIM = 32 # head_dim; n_heads = HIDDEN / HEAD_DIM = 3 +MLP_HIDDEN = int(HIDDEN / 3 * 8) # 256, matches ZImageTransformerBlock FeedForward +ADALN_OUT = 4 * HIDDEN # 384, matches Z-Image adaLN_modulation +N_LAYERS = 2 # main transformer blocks +N_REFINER = 1 # noise_refiner / context_refiner blocks + + +# pylint: disable=attribute-defined-outside-init +class _Holder(torch.nn.Module): + """Empty container module - we attach children dynamically.""" + + +def build_zimage_block(modulation: bool): + """Mirror ``ZImageTransformerBlock``'s diffusers-side module layout.""" + block = _Holder() + # Attention module (diffusers Attention class with custom processor) + block.attention = _Holder() + block.attention.to_q = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + block.attention.to_k = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + block.attention.to_v = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + # to_out is a ModuleList in diffusers; index 0 is the Linear, index 1 is Dropout + block.attention.to_out = torch.nn.ModuleList([ + torch.nn.Linear(HIDDEN, HIDDEN, bias=False), + torch.nn.Dropout(0.0), + ]) + # qk_norm RMSNorms - present when qk_norm=True + block.attention.norm_q = torch.nn.RMSNorm(HEAD_DIM) + block.attention.norm_k = torch.nn.RMSNorm(HEAD_DIM) + + # FeedForward (diffusers FeedForward class) + block.feed_forward = _Holder() + block.feed_forward.net = torch.nn.ModuleList() + proj_act = _Holder() # GELU activation wrapping a proj Linear + proj_act.proj = torch.nn.Linear(HIDDEN, MLP_HIDDEN, bias=True) + block.feed_forward.net.append(proj_act) + block.feed_forward.net.append(torch.nn.Dropout(0.0)) + block.feed_forward.net.append(torch.nn.Linear(MLP_HIDDEN, HIDDEN, bias=True)) + + # Per-block RMSNorms + block.attention_norm1 = torch.nn.RMSNorm(HIDDEN) + block.ffn_norm1 = torch.nn.RMSNorm(HIDDEN) + block.attention_norm2 = torch.nn.RMSNorm(HIDDEN) + block.ffn_norm2 = torch.nn.RMSNorm(HIDDEN) + + if modulation: + block.adaLN_modulation = torch.nn.Sequential( + torch.nn.Linear(HIDDEN, ADALN_OUT, bias=True), + ) + return block + + +def build_mock_transformer(): + """Build a torch.nn.Module mimicking ``ZImageTransformer2DModel``. + + Mirrors the paths real Z-Image LoRAs target: + ``layers.X.attention.{to_q,to_k,to_v,to_out.0}``, + ``layers.X.feed_forward.net.{0.proj,2}``, + ``layers.X.adaLN_modulation.0``, plus the four per-block RMSNorms and + refiner stacks (``noise_refiner``, ``context_refiner``). + """ + transformer = _Holder() + transformer.layers = torch.nn.ModuleList([build_zimage_block(modulation=False) for _ in range(N_LAYERS)]) + transformer.noise_refiner = torch.nn.ModuleList([build_zimage_block(modulation=True) for _ in range(N_REFINER)]) + transformer.context_refiner = torch.nn.ModuleList([build_zimage_block(modulation=True) for _ in range(N_REFINER)]) + return transformer + + +class _MockZImagePipeline: + """Class name carries 'ZImage' so name-based model-type dispatch routes correctly.""" + + def __init__(self, transformer): + self.transformer = transformer + self.text_encoder = None + + +class _MockZImageSdModel: + """Outer wrapper exposing pipe + network_layer_mapping for + ``lora_convert.assign_network_names_to_compvis_modules`` to write onto.""" + + def __init__(self, pipe): + self.pipe = pipe + self.network_layer_mapping = {} + self.embedding_db = None + self.__class__.__name__ = 'ZImagePipeline' # belt-and-suspenders + + +def install_mock_pipe(): + """Set shared.sd_model to a mock exposing a Z-Image-shaped transformer. + + Each test calls this fresh so any prior network_layer_name stamps don't + leak across tests. Writes directly to ``model_data.sd_model`` to bypass + the ModelData lock that no-ops the public setter outside webui startup. + """ + transformer = build_mock_transformer() + pipe = _MockZImagePipeline(transformer) + sd_model = _MockZImageSdModel(pipe) + from modules.modeldata import model_data + model_data.sd_model = sd_model + return sd_model + + +# ============================================================ +# State-dict synthesizers (one per family/format) +# ============================================================ + +RANK_LORA = 8 +RANK_LOKR = 4 +LOKR_W1_DIM = 8 + + +def sd_lora_bfl_to_q(): + """BFL / AI-toolkit LoRA on a split-QKV target. Mirrors 80sFantasyZBase.""" + return { + 'diffusion_model.layers.0.attention.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.0.attention.to_q.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_bfl_adaln(): + """BFL LoRA on adaLN_modulation (the modulation Linear in Sequential).""" + # Only refiner blocks have modulation in our mock; index path with refiner + return { + 'diffusion_model.noise_refiner.0.adaLN_modulation.0.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.noise_refiner.0.adaLN_modulation.0.lora_B.weight': torch.randn(ADALN_OUT, RANK_LORA), + } + + +def sd_lora_peft_to_out(): + """PEFT-format LoRA on attention.to_out.0. Mirrors FameGrid_Revolution_ZIB_BOLD.""" + return { + 'transformer.layers.1.attention.to_out.0.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'transformer.layers.1.attention.to_out.0.lora_up.weight': torch.randn(HIDDEN, RANK_LORA), + 'transformer.layers.1.attention.to_out.0.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_kohya_to_q(): + """Kohya-format LoRA on attention.to_q (flat underscore path).""" + return { + 'lora_unet_layers_0_attention_to_q.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'lora_unet_layers_0_attention_to_q.lora_up.weight': torch.randn(HIDDEN, RANK_LORA), + 'lora_unet_layers_0_attention_to_q.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_legacy_fused_qkv(): + """Legacy Z-Image fused ``attention.qkv`` (BFL prefix). Loader splits the up-weight.""" + # up has 3*HIDDEN rows (Q stacked over K stacked over V), down is shared + return { + 'diffusion_model.layers.0.attention.qkv.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.0.attention.qkv.lora_B.weight': torch.randn(3 * HIDDEN, RANK_LORA), + 'diffusion_model.layers.0.attention.qkv.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_legacy_attention_out_alias(): + """Legacy ``attention.out`` alias - loader renames to ``attention.to_out.0``.""" + return { + 'diffusion_model.layers.1.attention.out.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.1.attention.out.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_with_dora_scale(): + """LoRA carrying a dora_scale companion vector to exercise DoRA threading.""" + return { + 'transformer.layers.0.attention.to_v.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'transformer.layers.0.attention.to_v.lora_up.weight': torch.randn(HIDDEN, RANK_LORA), + 'transformer.layers.0.attention.to_v.dora_scale': torch.randn(HIDDEN), + } + + +def sd_lokr_bfl_adaln(): + """BFL LoKR on the modulation linear. Mirrors after_dark_2_zib_lokr's adaLN target.""" + # w1=(HIDDEN/LOKR_W1_DIM, LOKR_W1_DIM), w2=(LOKR_W1_DIM, ADALN_OUT*HIDDEN/(LOKR_W1_DIM*HIDDEN)) + # Pick a factorization that yields kron shape == (ADALN_OUT, HIDDEN) + return { + 'diffusion_model.noise_refiner.0.adaLN_modulation.0.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM), + 'diffusion_model.noise_refiner.0.adaLN_modulation.0.lokr_w2': torch.randn(ADALN_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM), + 'diffusion_model.noise_refiner.0.adaLN_modulation.0.alpha': torch.tensor(float(LOKR_W1_DIM)), + } + + +def sd_lokr_legacy_fused_qkv(): + """Legacy fused qkv LoKR - loader defers split to apply time via NetworkModuleLokrChunk.""" + # w1 small, w2 must yield 3*HIDDEN rows after kron; pick LOKR_W1_DIM and (3*HIDDEN/LOKR_W1_DIM, HIDDEN/LOKR_W1_DIM) + return { + 'diffusion_model.layers.0.attention.qkv.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM), + 'diffusion_model.layers.0.attention.qkv.lokr_w2': torch.randn((3 * HIDDEN) // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM), + 'diffusion_model.layers.0.attention.qkv.alpha': torch.tensor(float(LOKR_W1_DIM)), + } + + +def sd_loha_bfl_proj(): + """LoHA on attention.to_out.0 (non-fused; LoHA on fused qkv is skipped by the loader).""" + return { + 'diffusion_model.layers.1.attention.to_out.0.hada_w1_a': torch.randn(HIDDEN, RANK_LORA), + 'diffusion_model.layers.1.attention.to_out.0.hada_w1_b': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.1.attention.to_out.0.hada_w2_a': torch.randn(HIDDEN, RANK_LORA), + 'diffusion_model.layers.1.attention.to_out.0.hada_w2_b': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.1.attention.to_out.0.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_loha_legacy_fused_qkv_skipped(): + """LoHA on legacy fused qkv - loader skips with warning (no NetworkModuleHadaChunk for z-image's path).""" + return { + 'diffusion_model.layers.0.attention.qkv.hada_w1_a': torch.randn(3 * HIDDEN, RANK_LORA), + 'diffusion_model.layers.0.attention.qkv.hada_w1_b': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.0.attention.qkv.hada_w2_a': torch.randn(3 * HIDDEN, RANK_LORA), + 'diffusion_model.layers.0.attention.qkv.hada_w2_b': torch.randn(RANK_LORA, HIDDEN), + } + + +def sd_oft_lycoris_proj(): + """OFT LyCORIS-style oft_diag on attention.to_v (non-fused).""" + # OFT 3-D oft_blocks shape: (num_blocks, block_size, block_size). oft_diag is per-block. + num_blocks = 4 + block_size = HIDDEN // num_blocks + return { + 'diffusion_model.layers.0.attention.to_v.oft_blocks': torch.randn(num_blocks, block_size, block_size) * 0.01, + 'diffusion_model.layers.0.attention.to_v.oft_diag': torch.ones(num_blocks, block_size), + 'diffusion_model.layers.0.attention.to_v.alpha': torch.tensor(0.001), + } + + +def sd_oft_legacy_fused_qkv_skipped(): + """OFT on legacy fused qkv - loader skips with warning (no row-sliceable OFT structure).""" + num_blocks = 4 + block_size = (3 * HIDDEN) // num_blocks + return { + 'diffusion_model.layers.0.attention.qkv.oft_blocks': torch.randn(num_blocks, block_size, block_size) * 0.01, + } + + +# ============================================================ +# Helpers: write state dict to disk, mock NetworkOnDisk +# ============================================================ + + +class TempLora: + """Context manager: writes a state dict to a temp safetensors file and yields + a ``_MockNetworkOnDisk`` pointing at it. Cleans up on exit.""" + + def __init__(self, state_dict, name='test'): + self.state_dict = state_dict + self.name = name + self.path = None + + def __enter__(self): + # safetensors requires contiguous tensors + sd = {k: v.contiguous() if isinstance(v, torch.Tensor) else v for k, v in self.state_dict.items()} + fd, self.path = tempfile.mkstemp(suffix='.safetensors', prefix=f'{self.name}_') + os.close(fd) + safetensors.torch.save_file(sd, self.path) + return _MockNetworkOnDisk(self.path, self.name) + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.path and os.path.exists(self.path): + os.unlink(self.path) + + +class _MockNetworkOnDisk: + """Stand-in for ``network.NetworkOnDisk`` exposing only the attributes + the native loaders read.""" + + def __init__(self, filename, name): + self.filename = filename + self.name = name + self.shorthash = '' + self.sd_version = 'unknown' + + +def assert_shape(t: torch.Tensor, expected_shape, label=''): + actual = tuple(t.shape) + assert actual == tuple(expected_shape), f'{label}: shape {actual}, expected {tuple(expected_shape)}' + + +def make_network_for_module(net_module: network.NetworkModule, te_mul: float = 1.0, unet_mul: float = 1.0): + """Wire a NetworkModule's parent ``Network`` with given multipliers so + calc_updown has a meaningful multiplier()/calc_scale() context.""" + net_module.network.te_multiplier = te_mul + net_module.network.unet_multiplier = unet_mul + return net_module + + +# ============================================================ +# Tests - parsing primitives +# ============================================================ + +CAT_PARSE = category('parse') + + +def test_parse_key_all_prefixes(): + """parse_key recognizes BFL, PEFT, kohya, and bare keys with the right base+suffix.""" + cases = [ + # BFL prefix -> dotted base, lora_A normalized to lora_down + ('diffusion_model.layers.0.attention.to_q.lora_A.weight', + Z.LORA_SUFFIXES, + ('lora_transformer_layers_0_attention_to_q', 'lora_down.weight')), + # PEFT prefix + ('transformer.layers.0.attention.to_v.lora_B.weight', + Z.LORA_SUFFIXES, + ('lora_transformer_layers_0_attention_to_v', 'lora_up.weight')), + # kohya flat underscore form + ('lora_unet_layers_0_attention_to_k.lora_down.weight', + Z.LORA_SUFFIXES, + ('lora_transformer_layers_0_attention_to_k', 'lora_down.weight')), + # Bare dotted (no recognized prefix) + ('layers.0.attention.to_out.0.lora_A.weight', + Z.LORA_SUFFIXES, + ('lora_transformer_layers_0_attention_to_out_0', 'lora_down.weight')), + # Unrelated keys reject cleanly + ('random.unrelated.key', Z.LORA_SUFFIXES, None), + ] + for key, suffixes, expected in cases: + got = Z.parse_key(key, suffixes) + assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}' + return True + + +def test_marker_disambiguation(): + """Each family's markers reject other families' files.""" + pure_lora = { + 'lora_unet_layers_0_attention_to_q.lora_down.weight': torch.zeros(1, 1), + 'lora_unet_layers_0_attention_to_q.lora_up.weight': torch.zeros(1, 1), + } + assert Z.has_marker(pure_lora, Z.LORA_MARKERS) + assert not Z.has_marker(pure_lora, Z.LOKR_MARKERS) + assert not Z.has_marker(pure_lora, Z.LOHA_MARKERS) + assert not Z.has_marker(pure_lora, Z.OFT_MARKERS) + + pure_lokr = { + 'diffusion_model.layers.0.attention.to_q.lokr_w1': torch.zeros(1, 1), + 'diffusion_model.layers.0.attention.to_q.lokr_w2': torch.zeros(1, 1), + } + assert Z.has_marker(pure_lokr, Z.LOKR_MARKERS) + assert not Z.has_marker(pure_lokr, Z.LORA_MARKERS) + assert not Z.has_marker(pure_lokr, Z.LOHA_MARKERS) + assert not Z.has_marker(pure_lokr, Z.OFT_MARKERS) + return True + + +# ============================================================ +# Tests - loaders end-to-end +# ============================================================ + +CAT_LOADER = category('loader') + + +def _load_via(try_fn, state_dict, name='test'): + install_mock_pipe() + with TempLora(state_dict, name=name) as nod: + return try_fn(name, nod, lora_scale=1.0) + + +def test_lora_bfl_split_qkv(): + """BFL-format LoRA on attention.to_q binds correctly.""" + net = _load_via(Z.try_load_lora, sd_lora_bfl_to_q()) + assert net is not None and len(net.modules) == 1, f'expected 1 module, got {net.modules if net else None}' + assert 'lora_transformer_layers_0_attention_to_q' in net.modules + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_lora.NetworkModuleLora) + return True + + +def test_lora_peft_to_out(): + """PEFT-format LoRA on attention.to_out.0 binds correctly.""" + net = _load_via(Z.try_load_lora, sd_lora_peft_to_out()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_1_attention_to_out_0' in net.modules + return True + + +def test_lora_kohya_to_q(): + """Kohya flat-underscore LoRA on attention.to_q binds correctly.""" + net = _load_via(Z.try_load_lora, sd_lora_kohya_to_q()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_0_attention_to_q' in net.modules + return True + + +def test_lora_bfl_adaln(): + """BFL LoRA on adaLN_modulation.0 (the Linear inside Sequential).""" + net = _load_via(Z.try_load_lora, sd_lora_bfl_adaln()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_noise_refiner_0_adaLN_modulation_0' in net.modules + return True + + +def test_lora_legacy_fused_qkv_chunked(): + """Legacy fused attention.qkv is split into 3 chunks targeting to_q/to_k/to_v. + + The up-weight is chunked along dim 0 at load time; the down-weight is shared. + """ + net = _load_via(Z.try_load_lora, sd_lora_legacy_fused_qkv()) + assert net is not None and len(net.modules) == 3, f'expected 3 chunked modules, got {net.modules if net else None}' + expected = { + 'lora_transformer_layers_0_attention_to_q', + 'lora_transformer_layers_0_attention_to_k', + 'lora_transformer_layers_0_attention_to_v', + } + assert set(net.modules) == expected, f'got {set(net.modules)}' + # Each chunked module's up-weight is (HIDDEN, RANK), not (3*HIDDEN, RANK) + for nk, mod in net.modules.items(): + assert_shape(mod.up_model.weight, (HIDDEN, RANK_LORA), label=nk) + return True + + +def test_lora_legacy_attention_out_alias(): + """Legacy attention.out alias is renamed to attention.to_out.0.""" + net = _load_via(Z.try_load_lora, sd_lora_legacy_attention_out_alias()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_1_attention_to_out_0' in net.modules + return True + + +def test_lora_dora_threading(): + """dora_scale flows into NetworkModuleLora.dora_scale.""" + net = _load_via(Z.try_load_lora, sd_lora_with_dora_scale()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert mod.dora_scale is not None, 'dora_scale not threaded into NetworkModule' + return True + + +def test_lokr_bfl_adaln(): + """BFL LoKR on adaLN_modulation.0 binds via NetworkModuleLokr.""" + net = _load_via(Z.try_load_lokr, sd_lokr_bfl_adaln()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_noise_refiner_0_adaLN_modulation_0' in net.modules + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_lokr.NetworkModuleLokr) + return True + + +def test_lokr_legacy_fused_qkv_chunked(): + """Legacy fused LoKR on attention.qkv emits 3 chunked modules via NetworkModuleLokrChunk.""" + net = _load_via(Z.try_load_lokr, sd_lokr_legacy_fused_qkv()) + assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}' + expected = { + 'lora_transformer_layers_0_attention_to_q', + 'lora_transformer_layers_0_attention_to_k', + 'lora_transformer_layers_0_attention_to_v', + } + assert set(net.modules) == expected + for nk, mod in net.modules.items(): + assert isinstance(mod, network_lokr.NetworkModuleLokrChunk), f'{nk}: type={type(mod).__name__}' + return True + + +def test_loha_bfl_proj(): + """BFL LoHA on a non-fused proj target binds via NetworkModuleHada.""" + net = _load_via(Z.try_load_loha, sd_loha_bfl_proj()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_hada.NetworkModuleHada) + return True + + +def test_loha_legacy_fused_qkv_skipped(): + """LoHA on legacy fused attention.qkv is skipped with a warning (no chunk variant).""" + net = _load_via(Z.try_load_loha, sd_loha_legacy_fused_qkv_skipped()) + # Either no net returned (nothing matched) or net with zero modules + assert net is None or len(net.modules) == 0, f'expected no modules, got {net.modules if net else None}' + return True + + +def test_oft_lycoris_no_npe(): + """OFT loader handles LyCORIS oft_diag files without NoneType-attr errors. + + Regression check for the constraint guard in network_oft.py:58. + """ + net = _load_via(Z.try_load_oft, sd_oft_lycoris_proj()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_oft.NetworkModuleOFT) + return True + + +def test_oft_legacy_fused_qkv_skipped(): + """OFT on legacy fused attention.qkv is skipped (no row-sliceable structure).""" + net = _load_via(Z.try_load_oft, sd_oft_legacy_fused_qkv_skipped()) + assert net is None or len(net.modules) == 0, f'expected no modules, got {net.modules if net else None}' + return True + + +# ============================================================ +# Tests - calc_updown shape sanity +# ============================================================ + +CAT_MATH = category('math') + + +def test_lora_calc_updown_shape(): + """NetworkModuleLora.calc_updown emits the right shape against a target weight.""" + net = _load_via(Z.try_load_lora, sd_lora_bfl_to_q()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoRA calc_updown') + return True + + +def test_lokr_calc_updown_shape(): + """NetworkModuleLokr.calc_updown produces a tensor matching the target's shape.""" + net = _load_via(Z.try_load_lokr, sd_lokr_bfl_adaln()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(ADALN_OUT, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoKR calc_updown') + return True + + +def test_lokr_chunk_calc_updown_shape(): + """LokrChunk returns the designated row range; chunk shape matches the split target.""" + net = _load_via(Z.try_load_lokr, sd_lokr_legacy_fused_qkv()) + assert net is not None + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) # one Q/K/V projection's shape + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LokrChunk calc_updown') + return True + + +def test_loha_calc_updown_shape(): + """NetworkModuleHada produces shapes matching the target.""" + net = _load_via(Z.try_load_loha, sd_loha_bfl_proj()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoHA calc_updown') + return True + + +def test_oft_calc_updown_shape(): + """NetworkModuleOFT calc_updown shape sanity against the target weight.""" + net = _load_via(Z.try_load_oft, sd_oft_lycoris_proj()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='OFT calc_updown') + return True + + +# ============================================================ +# Test runner +# ============================================================ + + +def run_tests(): + t0 = time.time() + + log.warning('=== Parsing primitives ===') + for fn in [test_parse_key_all_prefixes, test_marker_disambiguation]: + run_test(CAT_PARSE, fn) + + log.warning('=== Loaders ===') + for fn in [ + test_lora_bfl_split_qkv, + test_lora_peft_to_out, + test_lora_kohya_to_q, + test_lora_bfl_adaln, + test_lora_legacy_fused_qkv_chunked, + test_lora_legacy_attention_out_alias, + test_lora_dora_threading, + test_lokr_bfl_adaln, + test_lokr_legacy_fused_qkv_chunked, + test_loha_bfl_proj, + test_loha_legacy_fused_qkv_skipped, + test_oft_lycoris_no_npe, + test_oft_legacy_fused_qkv_skipped, + ]: + run_test(CAT_LOADER, fn) + + log.warning('=== calc_updown shape sanity ===') + for fn in [ + test_lora_calc_updown_shape, + test_lokr_calc_updown_shape, + test_lokr_chunk_calc_updown_shape, + test_loha_calc_updown_shape, + test_oft_calc_updown_shape, + ]: + run_test(CAT_MATH, fn) + + elapsed = time.time() - t0 + log.warning('=== Results ===') + total_pass = 0 + total_fail = 0 + for cat, info in results.items(): + status = 'PASS' if info['failed'] == 0 else 'FAIL' + log.info(f' {cat}: {info["passed"]} passed, {info["failed"]} failed [{status}]') + total_pass += info['passed'] + total_fail += info['failed'] + log.warning(f'Total: {total_pass} passed, {total_fail} failed in {elapsed:.2f}s') + return total_fail == 0 + + +if __name__ == '__main__': + ok = run_tests() + sys.exit(0 if ok else 1) From be46362e3e5648b1eacac3b7daf0849973022bd2 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 22:43:05 +0100 Subject: [PATCH 07/13] fix(chroma): prepend lora_transformer_ to LoKR slice_info keys try_load_lokr renamed slice_info keys via static_rename but didn't add the lora_transformer_ prefix that apply_static_rename adds to the groups dict. slice_info.get(network_key) always returned None for fused targets, so the loader fell back to NetworkModuleLokr instead of NetworkModuleLokrSliceChunk on every fused-QKV / fused-linear1 adapter. The full kron(w1, w2) was applied against split target modules, either shape-mismatching at apply time or broadcasting wrong. Caught by test_lokr_bfl_img_attn_qkv_slice_chunked. No real-world chroma LoKR adapters on fused targets are known in the wild. --- pipelines/chroma/chroma_lora.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index 382bc6f3f..c92635948 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -191,7 +191,10 @@ def try_load_lokr(name, network_on_disk, lora_scale): groups = group_by_suffixes(state_dict, LOKR_SUFFIXES) groups, slice_info = expand_chroma_fused_lokr(groups) groups = apply_static_rename(groups, static_rename) - slice_info = {static_rename.get(k, k): v for k, v in slice_info.items()} + # Mirror apply_static_rename: the slice_info dict must use the same final + # network keys as ``groups`` (static rename applied + ``lora_transformer_`` + # prefix) so the per-target lookup in the loop matches. + slice_info = {'lora_transformer_' + static_rename.get(k, k): v for k, v in slice_info.items()} unmapped = 0 for network_key, w in groups.items(): From 1f309fe8364e07419f9e1791a4819c7161ea1f8a Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 22:43:05 +0100 Subject: [PATCH 08/13] test(chroma): offline tests for native adapter loaders Covers chroma's four-family surface plus the Flux-to-diffusers rename and the unique single-block linear1 unequal-chunk slicing. Mock transformer matches diffusers.ChromaTransformer2DModel (transformer_blocks with FluxAttention(added_kv_proj_dim), single transformer_blocks with pre-only attn + proj_mlp + proj_out, plus distilled_guidance_layer). Formats exercised: - BFL / AI-toolkit - kohya - PEFT - LyCORIS oft_diag install_mock_pipe patches chroma_lora.QKV_DIMS and LINEAR1_DIMS to the test scale (HIDDEN=96, MLP_HIDDEN=384); the module otherwise hardcodes Chroma1-HD's 3072 / 12288. --- test/test-chroma-native-adapters.py | 955 ++++++++++++++++++++++++++++ 1 file changed, 955 insertions(+) create mode 100644 test/test-chroma-native-adapters.py diff --git a/test/test-chroma-native-adapters.py b/test/test-chroma-native-adapters.py new file mode 100644 index 000000000..087b8db82 --- /dev/null +++ b/test/test-chroma-native-adapters.py @@ -0,0 +1,955 @@ +#!/usr/bin/env python +""" +Offline unit tests for Chroma native adapter loaders. + +Covers the four native families currently supported by ``pipelines.chroma.chroma_lora`` +(LoRA, LoKR, LoHA, OFT) plus DoRA threading via the universal +``NetworkModule.finalize_updown`` hook. + +Tests build a mock Chroma-shaped transformer, write synthetic safetensors +files for each adapter format observed in the wild, and exercise the full +loader path including the Flux-to-diffusers rename table and the unique +single-block ``linear1`` unequal-chunk slicing. + +Save formats are cross-referenced against real Chroma LoRAs: + +- BFL / AI-toolkit (``diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight``): + e.g. ``Chroma - Lenovo UltraReal`` +- kohya (``lora_unet_double_blocks_0_img_attn_proj.lora_down.weight``): + e.g. ``90s_anime_aesthetic_Chroma`` +- PEFT (``transformer.transformer_blocks.0.attn.to_q.lora_down.weight``) + +Chroma LoRAs are trained against the Flux block layout (``double_blocks``, +``single_blocks``) regardless of save format. The diffusers +``ChromaTransformer2DModel`` exposes split-attention modules at +``transformer_blocks.X.attn.{to_q,to_k,to_v,...}`` and +``single_transformer_blocks.X.{attn.*, proj_mlp, proj_out}``. The loader +path-rewrites Flux paths to diffusers names and handles two distinct +fused-weight layouts: + +- **Equal chunks** (double_blocks img_attn.qkv / txt_attn.qkv at + ``[HIDDEN, HIDDEN, HIDDEN]``): LoRA chunks at load via ``torch.chunk``; + LoKR defers via ``NetworkModuleLokrChunk``. +- **Unequal chunks** (single_blocks linear1 at + ``[HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN]``): LoRA slices row ranges at load; + LoKR defers via ``NetworkModuleLokrSliceChunk``. + +LoHA and OFT on fused targets are skipped with a warning (no slice variant). + +The ``distilled_guidance_layer`` (Chroma's central modulation generator that +replaces Flux's per-block ``norm1.linear``) is a real module path that +``assign_network_names_to_compvis_modules`` registers, so LoRAs targeting it +pass through unchanged. + +No running server required. + +Usage: + python test/test-chroma-native-adapters.py +""" + +import os +import sys +import tempfile +import time + +import torch +import safetensors.torch + +script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, script_dir) +os.chdir(script_dir) + +os.environ['SD_INSTALL_QUIET'] = '1' + +# Bootstrap cmd_args before any module that pulls in shared.py. +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +_orig_argv = sys.argv +sys.argv = [sys.argv[0]] +try: + modules.cmd_args.parse_args() +finally: + sys.argv = _orig_argv +installer.add_args(modules.cmd_args.parser) +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 # pylint: disable=wrong-import-position +from modules.lora import ( # pylint: disable=wrong-import-position + network, network_lora, network_lokr, network_hada, network_oft, +) +from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position +from pipelines.chroma import chroma_lora as C # pylint: disable=wrong-import-position + + +# ============================================================ +# Test infrastructure +# ============================================================ + +results: dict[str, dict] = {} + + +def category(name: str): + if name not in results: + results[name] = {'passed': 0, 'failed': 0, 'tests': []} + return name + + +def record(cat: str, passed: bool, name: str, detail: str = ''): + status = 'PASS' if passed else 'FAIL' + results[cat]['passed' if passed else 'failed'] += 1 + results[cat]['tests'].append((status, name)) + msg = f' {status}: {name}' + if detail: + msg += f' ({detail})' + if passed: + log.info(msg) + else: + log.error(msg) + + +def run_test(cat: str, fn): + name = fn.__name__ + try: + ok = fn() + if ok is False: + record(cat, False, name) + else: + record(cat, True, name) + except AssertionError as e: + record(cat, False, name, str(e)) + except Exception as e: # pylint: disable=broad-except + record(cat, False, name, f'exception: {e}') + import traceback + traceback.print_exc() + + +# ============================================================ +# Mock Chroma transformer +# ============================================================ +# Shape constants chosen to mirror ChromaTransformer2DModel proportions +# while keeping tensors small. Real Chroma1-HD: inner_dim=3072, +# mlp_hidden=12288. We use HIDDEN=96, MLP_HIDDEN=384 (4x), so the unequal +# single-block linear1 partition [HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN] = +# [96, 96, 96, 384] (analogous to real [3072, 3072, 3072, 12288]). + +HIDDEN = 96 +HEAD_DIM = 32 # N_HEADS = HIDDEN / HEAD_DIM = 3 +MLP_RATIO = 4 +MLP_HIDDEN = HIDDEN * MLP_RATIO # 384 +QKV_FUSED_OUT = 3 * HIDDEN # 288 (img_attn.qkv / txt_attn.qkv output dim) +LINEAR1_OUT = 3 * HIDDEN + MLP_HIDDEN # 672 (single block linear1 fused output) +LINEAR2_IN = HIDDEN + MLP_HIDDEN # 480 (single block proj_out input - attn out + mlp out concat) +N_DOUBLE = 2 +N_SINGLE = 2 + + +# pylint: disable=attribute-defined-outside-init +class _Holder(torch.nn.Module): + """Empty container module - we attach children dynamically.""" + + +def build_double_block(): + """Mirror ``ChromaTransformerBlock``'s diffusers-side module layout. + + Uses ``FluxAttention(added_kv_proj_dim=dim)`` so both img-side and + context-side QKV + output projections are present. + """ + block = _Holder() + # FluxAttention sub-modules + block.attn = _Holder() + block.attn.to_q = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.to_k = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.to_v = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.to_out = torch.nn.ModuleList([ + torch.nn.Linear(HIDDEN, HIDDEN, bias=True), + torch.nn.Dropout(0.0), + ]) + block.attn.add_q_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.add_k_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.add_v_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.to_add_out = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.norm_q = torch.nn.RMSNorm(HEAD_DIM) + block.attn.norm_k = torch.nn.RMSNorm(HEAD_DIM) + block.attn.norm_added_q = torch.nn.RMSNorm(HEAD_DIM) + block.attn.norm_added_k = torch.nn.RMSNorm(HEAD_DIM) + + # FeedForward modules: net = [GELU(proj=Linear), Dropout, Linear] + block.ff = _Holder() + block.ff.net = torch.nn.ModuleList() + proj_act = _Holder() + proj_act.proj = torch.nn.Linear(HIDDEN, MLP_HIDDEN, bias=True) + block.ff.net.append(proj_act) + block.ff.net.append(torch.nn.Dropout(0.0)) + block.ff.net.append(torch.nn.Linear(MLP_HIDDEN, HIDDEN, bias=True)) + + block.ff_context = _Holder() + block.ff_context.net = torch.nn.ModuleList() + proj_act_ctx = _Holder() + proj_act_ctx.proj = torch.nn.Linear(HIDDEN, MLP_HIDDEN, bias=True) + block.ff_context.net.append(proj_act_ctx) + block.ff_context.net.append(torch.nn.Dropout(0.0)) + block.ff_context.net.append(torch.nn.Linear(MLP_HIDDEN, HIDDEN, bias=True)) + + # norm1 / norm1_context / norm2 / norm2_context are AdaLayerNormZeroPruned + # or LayerNorm(elementwise_affine=False) - no learnable weight at the + # block-norm level, so we don't need LoRA-targetable norm modules here. + + return block + + +def build_single_block(): + """Mirror ``ChromaSingleTransformerBlock`` - has proj_mlp + attn (pre_only) + proj_out.""" + block = _Holder() + block.attn = _Holder() + block.attn.to_q = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.to_k = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.to_v = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + block.attn.norm_q = torch.nn.RMSNorm(HEAD_DIM) + block.attn.norm_k = torch.nn.RMSNorm(HEAD_DIM) + # pre_only=True so no to_out + block.proj_mlp = torch.nn.Linear(HIDDEN, MLP_HIDDEN, bias=True) + block.proj_out = torch.nn.Linear(LINEAR2_IN, HIDDEN, bias=True) + return block + + +def build_mock_transformer(): + """Build a torch.nn.Module mimicking ``ChromaTransformer2DModel``.""" + transformer = _Holder() + transformer.transformer_blocks = torch.nn.ModuleList([build_double_block() for _ in range(N_DOUBLE)]) + transformer.single_transformer_blocks = torch.nn.ModuleList([build_single_block() for _ in range(N_SINGLE)]) + # distilled_guidance_layer - Chroma's central modulation approximator + # Minimal stand-in: just one linear submodule the tests can target + transformer.distilled_guidance_layer = _Holder() + transformer.distilled_guidance_layer.in_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + return transformer + + +class _MockChromaPipeline: + """Class name carries 'Chroma' so name-based model-type dispatch routes correctly.""" + + def __init__(self, transformer): + self.transformer = transformer + self.text_encoder = None + + +class _MockChromaSdModel: + """Outer wrapper holding pipe + network_layer_mapping.""" + + def __init__(self, pipe): + self.pipe = pipe + self.network_layer_mapping = {} + self.embedding_db = None + self.__class__.__name__ = 'ChromaPipeline' + + +def install_mock_pipe(): + """Set shared.sd_model to a mock exposing a Chroma-shaped transformer. + + Each test re-installs so any prior network_layer_name stamps don't leak. + Writes directly to model_data.sd_model to bypass the ModelData lock. + + Also patches ``chroma_lora.QKV_DIMS`` and ``chroma_lora.LINEAR1_DIMS`` to + match the test mock's scaled-down ``HIDDEN`` / ``MLP_HIDDEN``. The module + hardcodes Chroma1-HD's 3072 / 12288, which mismatches small test tensors + and causes ``split_fused_lora_group``'s + ``up.shape[0] != sum(dims)`` gate to reject every fused fixture. + """ + transformer = build_mock_transformer() + pipe = _MockChromaPipeline(transformer) + sd_model = _MockChromaSdModel(pipe) + from modules.modeldata import model_data + model_data.sd_model = sd_model + + # chroma_lora's get_block_counts() reads transformer.config.num_layers / + # num_single_layers. Stamp that here so build_static_rename gets the right + # block counts for the test mock (defaults are 19/38 which our 2/2 mock doesn't have). + transformer.config = _ChromaConfig(num_layers=N_DOUBLE, num_single_layers=N_SINGLE) + + # Patch the hardcoded Chroma1-HD dims to the test scale. + C.QKV_DIMS = [HIDDEN, HIDDEN, HIDDEN] + C.LINEAR1_DIMS = [HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN] + return sd_model + + +class _ChromaConfig: + def __init__(self, num_layers, num_single_layers): + self.num_layers = num_layers + self.num_single_layers = num_single_layers + + +# ============================================================ +# State-dict synthesizers (one per family/format) +# ============================================================ + +RANK_LORA = 8 +LOKR_W1_DIM = 8 + + +def sd_lora_bfl_img_attn_proj(): + """BFL LoRA on double-block img_attn.proj. + + BFL path maps to diffusers ``transformer_blocks.0.attn.to_out.0`` via + ``DOUBLE_RENAME_TEMPLATES``. + """ + return { + 'diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.double_blocks.0.img_attn.proj.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + 'diffusion_model.double_blocks.0.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_bfl_img_attn_qkv_fused(): + """BFL LoRA on fused img_attn.qkv. Loader splits up-weight along dim 0 into Q/K/V.""" + return { + 'diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight': torch.randn(QKV_FUSED_OUT, RANK_LORA), + 'diffusion_model.double_blocks.0.img_attn.qkv.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_bfl_txt_attn_qkv_fused(): + """BFL LoRA on fused txt_attn.qkv. Loader emits 3 chunks to add_{q,k,v}_proj.""" + return { + 'diffusion_model.double_blocks.0.txt_attn.qkv.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.double_blocks.0.txt_attn.qkv.lora_B.weight': torch.randn(QKV_FUSED_OUT, RANK_LORA), + } + + +def sd_lora_bfl_img_mlp(): + """BFL LoRA on double-block img_mlp.0 and img_mlp.2. + + img_mlp.0 -> ff.net.0.proj, img_mlp.2 -> ff.net.2. + """ + return { + 'diffusion_model.double_blocks.1.img_mlp.0.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.double_blocks.1.img_mlp.0.lora_B.weight': torch.randn(MLP_HIDDEN, RANK_LORA), + 'diffusion_model.double_blocks.1.img_mlp.2.lora_A.weight': torch.randn(RANK_LORA, MLP_HIDDEN), + 'diffusion_model.double_blocks.1.img_mlp.2.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_bfl_txt_mlp(): + """BFL LoRA on double-block txt_mlp.0 and txt_mlp.2 - context side.""" + return { + 'diffusion_model.double_blocks.0.txt_mlp.0.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.double_blocks.0.txt_mlp.0.lora_B.weight': torch.randn(MLP_HIDDEN, RANK_LORA), + } + + +def sd_lora_bfl_single_linear1_unequal(): + """BFL LoRA on single-block linear1. + + linear1 fuses Q/K/V/proj_mlp at unequal dims [HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN]. + Loader emits 4 targets with unequal row-range chunks. + """ + return { + 'diffusion_model.single_blocks.0.linear1.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.single_blocks.0.linear1.lora_B.weight': torch.randn(LINEAR1_OUT, RANK_LORA), + } + + +def sd_lora_bfl_single_linear2(): + """BFL LoRA on single-block linear2 (-> single_transformer_blocks.X.proj_out).""" + return { + 'diffusion_model.single_blocks.0.linear2.lora_A.weight': torch.randn(RANK_LORA, LINEAR2_IN), + 'diffusion_model.single_blocks.0.linear2.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_kohya_img_attn_proj(): + """Kohya flat-underscore LoRA on img_attn.proj. Mirrors 90s_anime_aesthetic_Chroma.""" + return { + 'lora_unet_double_blocks_0_img_attn_proj.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'lora_unet_double_blocks_0_img_attn_proj.lora_up.weight': torch.randn(HIDDEN, RANK_LORA), + 'lora_unet_double_blocks_0_img_attn_proj.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_kohya_img_attn_qkv_fused(): + """Kohya LoRA on fused img_attn.qkv.""" + return { + 'lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight': torch.randn(QKV_FUSED_OUT, RANK_LORA), + 'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_peft_to_q(): + """PEFT-format LoRA targeting a split diffusers path (no rename, no chunking).""" + return { + 'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_distilled_guidance(): + """LoRA targeting Chroma's distilled_guidance_layer (passes through unchanged).""" + return { + 'diffusion_model.distilled_guidance_layer.in_proj.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.distilled_guidance_layer.in_proj.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_with_dora_scale(): + """LoRA with dora_scale companion to exercise DoRA threading.""" + return { + 'transformer.transformer_blocks.0.attn.to_v.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'transformer.transformer_blocks.0.attn.to_v.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + 'transformer.transformer_blocks.0.attn.to_v.dora_scale': torch.randn(HIDDEN), + } + + +def sd_lokr_bfl_img_attn_proj(): + """BFL LoKR on a non-fused proj target. Loader uses NetworkModuleLokr (no chunk).""" + return { + 'diffusion_model.double_blocks.0.img_attn.proj.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM), + 'diffusion_model.double_blocks.0.img_attn.proj.lokr_w2': torch.randn(HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM), + 'diffusion_model.double_blocks.0.img_attn.proj.alpha': torch.tensor(float(LOKR_W1_DIM)), + } + + +def sd_lokr_bfl_img_attn_qkv_equal_chunks(): + """BFL LoKR on fused img_attn.qkv (equal chunks). + + Loader emits 3 NetworkModuleLokrSliceChunk via row ranges [0:HIDDEN], [HIDDEN:2*HIDDEN], [2*HIDDEN:3*HIDDEN]. + (Chroma's implementation slices even equal-chunks via row ranges since the same logic handles both.) + """ + return { + 'diffusion_model.double_blocks.0.img_attn.qkv.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM), + 'diffusion_model.double_blocks.0.img_attn.qkv.lokr_w2': torch.randn(QKV_FUSED_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM), + 'diffusion_model.double_blocks.0.img_attn.qkv.alpha': torch.tensor(float(LOKR_W1_DIM)), + } + + +def sd_lokr_bfl_single_linear1_unequal(): + """BFL LoKR on fused single-block linear1 (UNEQUAL chunks). + + Loader emits 4 NetworkModuleLokrSliceChunk with row ranges matching + [HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN] partitions. + """ + return { + 'diffusion_model.single_blocks.0.linear1.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM), + 'diffusion_model.single_blocks.0.linear1.lokr_w2': torch.randn(LINEAR1_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM), + 'diffusion_model.single_blocks.0.linear1.alpha': torch.tensor(float(LOKR_W1_DIM)), + } + + +def sd_loha_bfl_img_attn_proj(): + """LoHA on a non-fused target binds via NetworkModuleHada.""" + return { + 'diffusion_model.double_blocks.1.img_attn.proj.hada_w1_a': torch.randn(HIDDEN, RANK_LORA), + 'diffusion_model.double_blocks.1.img_attn.proj.hada_w1_b': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.double_blocks.1.img_attn.proj.hada_w2_a': torch.randn(HIDDEN, RANK_LORA), + 'diffusion_model.double_blocks.1.img_attn.proj.hada_w2_b': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_loha_bfl_img_attn_qkv_skipped(): + """LoHA on fused img_attn.qkv is dropped by the loader (no slice variant for LoHA).""" + return { + 'diffusion_model.double_blocks.0.img_attn.qkv.hada_w1_a': torch.randn(QKV_FUSED_OUT, RANK_LORA), + 'diffusion_model.double_blocks.0.img_attn.qkv.hada_w1_b': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.double_blocks.0.img_attn.qkv.hada_w2_a': torch.randn(QKV_FUSED_OUT, RANK_LORA), + 'diffusion_model.double_blocks.0.img_attn.qkv.hada_w2_b': torch.randn(RANK_LORA, HIDDEN), + } + + +def sd_oft_bfl_img_attn_proj(): + """OFT (LyCORIS oft_diag form) on non-fused target.""" + num_blocks = 4 + block_size = HIDDEN // num_blocks + return { + 'diffusion_model.double_blocks.0.img_attn.proj.oft_blocks': torch.randn(num_blocks, block_size, block_size) * 0.01, + 'diffusion_model.double_blocks.0.img_attn.proj.oft_diag': torch.ones(num_blocks, block_size), + 'diffusion_model.double_blocks.0.img_attn.proj.alpha': torch.tensor(0.001), + } + + +def sd_oft_bfl_img_attn_qkv_skipped(): + """OFT on fused img_attn.qkv - dropped by the loader (OFT structure tied to out_features).""" + num_blocks = 4 + block_size = QKV_FUSED_OUT // num_blocks + return { + 'diffusion_model.double_blocks.0.img_attn.qkv.oft_blocks': torch.randn(num_blocks, block_size, block_size) * 0.01, + } + + +# ============================================================ +# Helpers +# ============================================================ + + +class TempLora: + """Context manager: writes a state dict to a temp safetensors file.""" + + def __init__(self, state_dict, name='test'): + self.state_dict = state_dict + self.name = name + self.path = None + + def __enter__(self): + sd = {k: v.contiguous() if isinstance(v, torch.Tensor) else v for k, v in self.state_dict.items()} + fd, self.path = tempfile.mkstemp(suffix='.safetensors', prefix=f'{self.name}_') + os.close(fd) + safetensors.torch.save_file(sd, self.path) + return _MockNetworkOnDisk(self.path, self.name) + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.path and os.path.exists(self.path): + os.unlink(self.path) + + +class _MockNetworkOnDisk: + def __init__(self, filename, name): + self.filename = filename + self.name = name + self.shorthash = '' + self.sd_version = 'unknown' + + +def assert_shape(t: torch.Tensor, expected_shape, label=''): + actual = tuple(t.shape) + assert actual == tuple(expected_shape), f'{label}: shape {actual}, expected {tuple(expected_shape)}' + + +def make_network_for_module(net_module: network.NetworkModule, te_mul: float = 1.0, unet_mul: float = 1.0): + net_module.network.te_multiplier = te_mul + net_module.network.unet_multiplier = unet_mul + return net_module + + +# ============================================================ +# Tests - parsing primitives +# ============================================================ + +CAT_PARSE = category('parse') + + +def test_parse_key_all_prefixes(): + """parse_key recognizes BFL, PEFT, kohya keys with the right flat_key + suffix. + + Chroma's parse_key returns (flat_key, suffix) where flat_key is the + underscored Flux-layout path (before static rename). + """ + cases = [ + # BFL prefix -> dotted base flattened to underscores + ('diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight', + C.LORA_SUFFIXES, + ('double_blocks_0_img_attn_proj', 'lora_down.weight')), + # PEFT prefix -> stays in diffusers form + ('transformer.transformer_blocks.0.attn.to_q.lora_B.weight', + C.LORA_SUFFIXES, + ('transformer_blocks_0_attn_to_q', 'lora_up.weight')), + # kohya prefix - already flat underscores + ('lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight', + C.LORA_SUFFIXES, + ('double_blocks_0_img_attn_qkv', 'lora_down.weight')), + # Unrelated keys reject cleanly + ('random.unrelated.key', C.LORA_SUFFIXES, None), + ] + for key, suffixes, expected in cases: + got = C.parse_key(key, suffixes) + assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}' + return True + + +def test_marker_disambiguation(): + """Each family's markers reject other families' files.""" + pure_lora = { + 'lora_unet_double_blocks_0_img_attn_proj.lora_down.weight': torch.zeros(1, 1), + 'lora_unet_double_blocks_0_img_attn_proj.lora_up.weight': torch.zeros(1, 1), + } + assert C.has_marker(pure_lora, C.LORA_MARKERS) + assert not C.has_marker(pure_lora, C.LOKR_MARKERS) + assert not C.has_marker(pure_lora, C.LOHA_MARKERS) + assert not C.has_marker(pure_lora, C.OFT_MARKERS) + + pure_lokr = { + 'diffusion_model.double_blocks.0.img_attn.proj.lokr_w1': torch.zeros(1, 1), + 'diffusion_model.double_blocks.0.img_attn.proj.lokr_w2': torch.zeros(1, 1), + } + assert C.has_marker(pure_lokr, C.LOKR_MARKERS) + assert not C.has_marker(pure_lokr, C.LORA_MARKERS) + return True + + +def test_static_rename_table(): + """build_static_rename produces correct Flux to diffusers path remappings. + + Verifies the rename templates against the documented chroma_lora layout: + double_blocks_X_img_attn_proj -> transformer_blocks_X_attn_to_out_0, etc. + """ + rename = C.build_static_rename(N_DOUBLE, N_SINGLE) + # Spot-check key remappings + assert rename['double_blocks_0_img_attn_proj'] == 'transformer_blocks_0_attn_to_out_0' + assert rename['double_blocks_0_txt_attn_proj'] == 'transformer_blocks_0_attn_to_add_out' + assert rename['double_blocks_1_img_mlp_0'] == 'transformer_blocks_1_ff_net_0_proj' + assert rename['double_blocks_1_img_mlp_2'] == 'transformer_blocks_1_ff_net_2' + assert rename['double_blocks_0_txt_mlp_0'] == 'transformer_blocks_0_ff_context_net_0_proj' + assert rename['single_blocks_0_linear2'] == 'single_transformer_blocks_0_proj_out' + return True + + +# ============================================================ +# Tests - loaders end-to-end +# ============================================================ + +CAT_LOADER = category('loader') + + +def _load_via(try_fn, state_dict, name='test'): + install_mock_pipe() + with TempLora(state_dict, name=name) as nod: + return try_fn(name, nod, lora_scale=1.0) + + +def test_lora_bfl_img_attn_proj(): + """BFL LoRA on img_attn.proj renames to attn.to_out.0.""" + net = _load_via(C.try_load_lora, sd_lora_bfl_img_attn_proj()) + assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}' + assert 'lora_transformer_transformer_blocks_0_attn_to_out_0' in net.modules + return True + + +def test_lora_bfl_img_attn_qkv_chunked(): + """BFL LoRA on fused img_attn.qkv emits 3 chunks targeting to_q/to_k/to_v.""" + net = _load_via(C.try_load_lora, sd_lora_bfl_img_attn_qkv_fused()) + assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}' + expected = { + 'lora_transformer_transformer_blocks_0_attn_to_q', + 'lora_transformer_transformer_blocks_0_attn_to_k', + 'lora_transformer_transformer_blocks_0_attn_to_v', + } + assert set(net.modules) == expected + # Each chunked up tensor has shape (HIDDEN, RANK), not (QKV_FUSED_OUT, RANK) + for nk, mod in net.modules.items(): + assert_shape(mod.up_model.weight, (HIDDEN, RANK_LORA), label=nk) + return True + + +def test_lora_bfl_txt_attn_qkv_chunked(): + """BFL LoRA on fused txt_attn.qkv emits 3 chunks targeting add_q/k/v_proj (context side).""" + net = _load_via(C.try_load_lora, sd_lora_bfl_txt_attn_qkv_fused()) + assert net is not None and len(net.modules) == 3 + expected = { + 'lora_transformer_transformer_blocks_0_attn_add_q_proj', + 'lora_transformer_transformer_blocks_0_attn_add_k_proj', + 'lora_transformer_transformer_blocks_0_attn_add_v_proj', + } + assert set(net.modules) == expected + return True + + +def test_lora_bfl_img_mlp(): + """img_mlp.0 -> ff.net.0.proj, img_mlp.2 -> ff.net.2.""" + net = _load_via(C.try_load_lora, sd_lora_bfl_img_mlp()) + assert net is not None and len(net.modules) == 2 + assert 'lora_transformer_transformer_blocks_1_ff_net_0_proj' in net.modules + assert 'lora_transformer_transformer_blocks_1_ff_net_2' in net.modules + return True + + +def test_lora_bfl_txt_mlp(): + """txt_mlp.0 -> ff_context.net.0.proj.""" + net = _load_via(C.try_load_lora, sd_lora_bfl_txt_mlp()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_transformer_blocks_0_ff_context_net_0_proj' in net.modules + return True + + +def test_lora_bfl_single_linear1_unequal_chunks(): + """BFL LoRA on single linear1 emits 4 targets with UNEQUAL row ranges. + + Partitions: [HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN] -> to_q, to_k, to_v, proj_mlp. + The first three chunks have (HIDDEN, RANK) up-shape; the fourth has (MLP_HIDDEN, RANK). + """ + net = _load_via(C.try_load_lora, sd_lora_bfl_single_linear1_unequal()) + assert net is not None and len(net.modules) == 4, f'got {net.modules if net else None}' + expected = { + 'lora_transformer_single_transformer_blocks_0_attn_to_q', + 'lora_transformer_single_transformer_blocks_0_attn_to_k', + 'lora_transformer_single_transformer_blocks_0_attn_to_v', + 'lora_transformer_single_transformer_blocks_0_proj_mlp', + } + assert set(net.modules) == expected + # proj_mlp has the MLP_HIDDEN chunk; QKV targets have HIDDEN + for nk, mod in net.modules.items(): + if nk.endswith('proj_mlp'): + assert_shape(mod.up_model.weight, (MLP_HIDDEN, RANK_LORA), label=nk) + else: + assert_shape(mod.up_model.weight, (HIDDEN, RANK_LORA), label=nk) + return True + + +def test_lora_bfl_single_linear2(): + """linear2 -> single_transformer_blocks.X.proj_out (no chunking).""" + net = _load_via(C.try_load_lora, sd_lora_bfl_single_linear2()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_single_transformer_blocks_0_proj_out' in net.modules + return True + + +def test_lora_kohya_img_attn_proj(): + """Kohya flat-underscore on non-fused target binds with same diffusers-path key as BFL.""" + net = _load_via(C.try_load_lora, sd_lora_kohya_img_attn_proj()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_transformer_blocks_0_attn_to_out_0' in net.modules + return True + + +def test_lora_kohya_img_attn_qkv_chunked(): + """Kohya fused img_attn.qkv splits into 3 chunks same as BFL form.""" + net = _load_via(C.try_load_lora, sd_lora_kohya_img_attn_qkv_fused()) + assert net is not None and len(net.modules) == 3 + expected = { + 'lora_transformer_transformer_blocks_0_attn_to_q', + 'lora_transformer_transformer_blocks_0_attn_to_k', + 'lora_transformer_transformer_blocks_0_attn_to_v', + } + assert set(net.modules) == expected + return True + + +def test_lora_peft_to_q(): + """PEFT format with diffusers paths passes through unchanged.""" + net = _load_via(C.try_load_lora, sd_lora_peft_to_q()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_transformer_blocks_0_attn_to_q' in net.modules + return True + + +def test_lora_distilled_guidance(): + """LoRA on distilled_guidance_layer passes through unchanged (real module path).""" + net = _load_via(C.try_load_lora, sd_lora_distilled_guidance()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_distilled_guidance_layer_in_proj' in net.modules + return True + + +def test_lora_dora_threading(): + """dora_scale flows into NetworkModuleLora.""" + net = _load_via(C.try_load_lora, sd_lora_with_dora_scale()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert mod.dora_scale is not None + return True + + +def test_lokr_bfl_img_attn_proj(): + """BFL LoKR on non-fused proj binds via NetworkModuleLokr (no chunk class).""" + net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_proj()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_transformer_blocks_0_attn_to_out_0' in net.modules + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_lokr.NetworkModuleLokr) and not isinstance(mod, network_lokr.NetworkModuleLokrChunk) + return True + + +def test_lokr_bfl_img_attn_qkv_slice_chunked(): + """BFL LoKR on fused img_attn.qkv emits 3 SliceChunk modules with row ranges. + + Chroma's expand_chroma_fused_lokr uses NetworkModuleLokrSliceChunk for all + fused targets (even equal-chunks), since the implementation is general + and start/end ranges express both forms. + """ + net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_qkv_equal_chunks()) + assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}' + expected = { + 'lora_transformer_transformer_blocks_0_attn_to_q', + 'lora_transformer_transformer_blocks_0_attn_to_k', + 'lora_transformer_transformer_blocks_0_attn_to_v', + } + assert set(net.modules) == expected + for nk, mod in net.modules.items(): + assert isinstance(mod, network_lokr.NetworkModuleLokrSliceChunk), f'{nk}: type={type(mod).__name__}' + # Each chunk row range is HIDDEN rows wide + assert mod.end_row - mod.start_row == HIDDEN, f'{nk}: range={mod.start_row}:{mod.end_row}' + return True + + +def test_lokr_bfl_single_linear1_unequal_chunks(): + """BFL LoKR on fused linear1 emits 4 SliceChunks with UNEQUAL ranges. + + Critical chroma-specific path: HIDDEN/HIDDEN/HIDDEN/MLP_HIDDEN partition. + """ + net = _load_via(C.try_load_lokr, sd_lokr_bfl_single_linear1_unequal()) + assert net is not None and len(net.modules) == 4, f'got {net.modules if net else None}' + # Check the proj_mlp chunk has the longer row range (MLP_HIDDEN) + proj_mlp_key = 'lora_transformer_single_transformer_blocks_0_proj_mlp' + assert proj_mlp_key in net.modules + proj_mlp = net.modules[proj_mlp_key] + assert proj_mlp.end_row - proj_mlp.start_row == MLP_HIDDEN, \ + f'proj_mlp range={proj_mlp.start_row}:{proj_mlp.end_row}, expected width={MLP_HIDDEN}' + # The three QKV chunks should each be HIDDEN rows wide + for proj in ('attn_to_q', 'attn_to_k', 'attn_to_v'): + nk = f'lora_transformer_single_transformer_blocks_0_{proj}' + mod = net.modules[nk] + assert mod.end_row - mod.start_row == HIDDEN, f'{nk}: range={mod.start_row}:{mod.end_row}' + return True + + +def test_loha_bfl_img_attn_proj(): + """LoHA on non-fused target binds via NetworkModuleHada.""" + net = _load_via(C.try_load_loha, sd_loha_bfl_img_attn_proj()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_hada.NetworkModuleHada) + return True + + +def test_loha_bfl_img_attn_qkv_skipped(): + """LoHA on fused img_attn.qkv is dropped (chroma has no LoHA chunk variant).""" + net = _load_via(C.try_load_loha, sd_loha_bfl_img_attn_qkv_skipped()) + assert net is None or len(net.modules) == 0 + return True + + +def test_oft_bfl_img_attn_proj(): + """LyCORIS oft_diag form loads on non-fused target without NoneType errors.""" + net = _load_via(C.try_load_oft, sd_oft_bfl_img_attn_proj()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_oft.NetworkModuleOFT) + return True + + +def test_oft_bfl_img_attn_qkv_skipped(): + """OFT on fused img_attn.qkv is dropped (no row-sliceable OFT structure).""" + net = _load_via(C.try_load_oft, sd_oft_bfl_img_attn_qkv_skipped()) + assert net is None or len(net.modules) == 0 + return True + + +# ============================================================ +# Tests - calc_updown shape sanity +# ============================================================ + +CAT_MATH = category('math') + + +def test_lora_calc_updown_shape(): + net = _load_via(C.try_load_lora, sd_lora_bfl_img_attn_proj()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoRA calc_updown') + return True + + +def test_lokr_calc_updown_shape(): + net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_proj()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoKR calc_updown') + return True + + +def test_lokr_slicechunk_equal_calc_updown_shape(): + """LokrSliceChunk with equal-width range produces (HIDDEN, HIDDEN) output.""" + net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_qkv_equal_chunks()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LokrSliceChunk equal range') + return True + + +def test_lokr_slicechunk_unequal_calc_updown_shape(): + """LokrSliceChunk on the proj_mlp chunk produces (MLP_HIDDEN, HIDDEN) output. + + This exercises the path that motivated NetworkModuleLokrSliceChunk's + existence: unequal partition where torch.chunk would not work. + """ + net = _load_via(C.try_load_lokr, sd_lokr_bfl_single_linear1_unequal()) + proj_mlp_key = 'lora_transformer_single_transformer_blocks_0_proj_mlp' + mod = make_network_for_module(net.modules[proj_mlp_key]) + target = torch.randn(MLP_HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LokrSliceChunk unequal proj_mlp') + return True + + +def test_loha_calc_updown_shape(): + net = _load_via(C.try_load_loha, sd_loha_bfl_img_attn_proj()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoHA calc_updown') + return True + + +def test_oft_calc_updown_shape(): + net = _load_via(C.try_load_oft, sd_oft_bfl_img_attn_proj()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='OFT calc_updown') + return True + + +# ============================================================ +# Test runner +# ============================================================ + + +def run_tests(): + t0 = time.time() + + log.warning('=== Parsing primitives ===') + for fn in [test_parse_key_all_prefixes, test_marker_disambiguation, test_static_rename_table]: + run_test(CAT_PARSE, fn) + + log.warning('=== Loaders ===') + for fn in [ + test_lora_bfl_img_attn_proj, + test_lora_bfl_img_attn_qkv_chunked, + test_lora_bfl_txt_attn_qkv_chunked, + test_lora_bfl_img_mlp, + test_lora_bfl_txt_mlp, + test_lora_bfl_single_linear1_unequal_chunks, + test_lora_bfl_single_linear2, + test_lora_kohya_img_attn_proj, + test_lora_kohya_img_attn_qkv_chunked, + test_lora_peft_to_q, + test_lora_distilled_guidance, + test_lora_dora_threading, + test_lokr_bfl_img_attn_proj, + test_lokr_bfl_img_attn_qkv_slice_chunked, + test_lokr_bfl_single_linear1_unequal_chunks, + test_loha_bfl_img_attn_proj, + test_loha_bfl_img_attn_qkv_skipped, + test_oft_bfl_img_attn_proj, + test_oft_bfl_img_attn_qkv_skipped, + ]: + run_test(CAT_LOADER, fn) + + log.warning('=== calc_updown shape sanity ===') + for fn in [ + test_lora_calc_updown_shape, + test_lokr_calc_updown_shape, + test_lokr_slicechunk_equal_calc_updown_shape, + test_lokr_slicechunk_unequal_calc_updown_shape, + test_loha_calc_updown_shape, + test_oft_calc_updown_shape, + ]: + run_test(CAT_MATH, fn) + + elapsed = time.time() - t0 + log.warning('=== Results ===') + total_pass = 0 + total_fail = 0 + for cat, info in results.items(): + status = 'PASS' if info['failed'] == 0 else 'FAIL' + log.info(f' {cat}: {info["passed"]} passed, {info["failed"]} failed [{status}]') + total_pass += info['passed'] + total_fail += info['failed'] + log.warning(f'Total: {total_pass} passed, {total_fail} failed in {elapsed:.2f}s') + return total_fail == 0 + + +if __name__ == '__main__': + ok = run_tests() + sys.exit(0 if ok else 1) From c7e8e7a029a0daf558a361496423e1a9a6a8fb66 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 23:00:12 +0100 Subject: [PATCH 09/13] test(ernie): offline tests for native adapter loaders Covers ernie's four-family surface (LoRA, LoKR, LoHA, OFT). ErnieImageAttention has fully split to_q / to_k / to_v with no fused QKV, and ErnieImageFeedForward has three separate Linear modules. The loader has no chunking, renames, or fused-target dispatch. Mock matches diffusers.ErnieImageTransformer2DModel (ErnieImageSharedAdaLNBlock with self_attention + mlp + RMSNorms, plus module-level adaLN_modulation Sequential and final_linear). Formats exercised: - BFL / AI-toolkit - kohya - BFL LoKR --- test/test-ernie-native-adapters.py | 636 +++++++++++++++++++++++++++++ 1 file changed, 636 insertions(+) create mode 100644 test/test-ernie-native-adapters.py diff --git a/test/test-ernie-native-adapters.py b/test/test-ernie-native-adapters.py new file mode 100644 index 000000000..8dd7fd2f5 --- /dev/null +++ b/test/test-ernie-native-adapters.py @@ -0,0 +1,636 @@ +#!/usr/bin/env python +""" +Offline unit tests for ERNIE-Image native adapter loaders. + +Covers the four native families currently supported by ``pipelines.ernie.ernie_lora`` +(LoRA, LoKR, LoHA, OFT) plus DoRA threading via the universal +``NetworkModule.finalize_updown`` hook. + +ERNIE-Image is the simplest native arch among z-image / chroma / ernie / flux2: +``ErnieImageAttention`` has fully split ``to_q`` / ``to_k`` / ``to_v`` Linear +modules (no fused QKV layout), and ``ErnieImageFeedForward`` exposes three +separate Linear modules (``gate_proj``, ``up_proj``, ``linear_fc2``). The +loader has no chunking, no renames, no fused-target dispatch - just direct +path-to-network-key conversion. + +Save formats are cross-referenced against real ERNIE-Image LoRAs: + +- BFL / AI-toolkit (``diffusion_model.layers.16.mlp.gate_proj.lora_A.weight``): + e.g. ``Ernie-Breast-Slider-v1`` +- kohya (``lora_unet_layers_0_mlp_gate_proj.lora_down.weight``): + e.g. ``ernie_image_radiancechromevoluptuous`` +- BFL LoKR (``diffusion_model.layers.0.mlp.gate_proj.lokr_w1``): + e.g. ``ERNIE_Anatomy_Male`` + +The diffusers ``ErnieImageTransformer2DModel`` layout +(``layers[i].self_attention.{to_q,to_k,to_v,to_out.0}``, +``layers[i].mlp.{gate_proj,up_proj,linear_fc2}``, plus the module-level +``adaLN_modulation.1`` Linear inside a Sequential, and top-level +``final_norm`` / ``final_linear``) is taken straight from +``diffusers.ErnieImageTransformer2DModel`` / +``ErnieImageSharedAdaLNBlock``. + +No running server required. + +Usage: + python test/test-ernie-native-adapters.py +""" + +import os +import sys +import tempfile +import time + +import torch +import safetensors.torch + +script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, script_dir) +os.chdir(script_dir) + +os.environ['SD_INSTALL_QUIET'] = '1' + +# Bootstrap cmd_args before any module that pulls in shared.py. +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +_orig_argv = sys.argv +sys.argv = [sys.argv[0]] +try: + modules.cmd_args.parse_args() +finally: + sys.argv = _orig_argv +installer.add_args(modules.cmd_args.parser) +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 # pylint: disable=wrong-import-position +from modules.lora import ( # pylint: disable=wrong-import-position + network, network_lora, network_lokr, network_hada, network_oft, +) +from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position +from pipelines.ernie import ernie_lora as E # pylint: disable=wrong-import-position + + +# ============================================================ +# Test infrastructure +# ============================================================ + +results: dict[str, dict] = {} + + +def category(name: str): + if name not in results: + results[name] = {'passed': 0, 'failed': 0, 'tests': []} + return name + + +def record(cat: str, passed: bool, name: str, detail: str = ''): + status = 'PASS' if passed else 'FAIL' + results[cat]['passed' if passed else 'failed'] += 1 + results[cat]['tests'].append((status, name)) + msg = f' {status}: {name}' + if detail: + msg += f' ({detail})' + if passed: + log.info(msg) + else: + log.error(msg) + + +def run_test(cat: str, fn): + name = fn.__name__ + try: + ok = fn() + if ok is False: + record(cat, False, name) + else: + record(cat, True, name) + except AssertionError as e: + record(cat, False, name, str(e)) + except Exception as e: # pylint: disable=broad-except + record(cat, False, name, f'exception: {e}') + import traceback + traceback.print_exc() + + +# ============================================================ +# Mock ERNIE-Image transformer +# ============================================================ +# ErnieImage upstream: hidden_size=3072, num_attention_heads=24, +# head_dim=128, ffn_hidden_size=8192. Test scale keeps proportions: +# HIDDEN=96, N_HEADS=3, HEAD_DIM=32, FFN_HIDDEN=256, ADALN_OUT=6*HIDDEN=576. + +HIDDEN = 96 +HEAD_DIM = 32 +FFN_HIDDEN = 256 +ADALN_OUT = 6 * HIDDEN +N_LAYERS = 2 + + +# pylint: disable=attribute-defined-outside-init +class _Holder(torch.nn.Module): + """Empty container module - we attach children dynamically.""" + + +def build_ernie_block(): + """Mirror ``ErnieImageSharedAdaLNBlock`` (single block type, no variants).""" + block = _Holder() + + block.self_attention = _Holder() + block.self_attention.to_q = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + block.self_attention.to_k = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + block.self_attention.to_v = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + block.self_attention.to_out = torch.nn.ModuleList([ + torch.nn.Linear(HIDDEN, HIDDEN, bias=False), + torch.nn.Dropout(0.0), + ]) + block.self_attention.norm_q = torch.nn.RMSNorm(HEAD_DIM) + block.self_attention.norm_k = torch.nn.RMSNorm(HEAD_DIM) + + block.mlp = _Holder() + block.mlp.gate_proj = torch.nn.Linear(HIDDEN, FFN_HIDDEN, bias=False) + block.mlp.up_proj = torch.nn.Linear(HIDDEN, FFN_HIDDEN, bias=False) + block.mlp.linear_fc2 = torch.nn.Linear(FFN_HIDDEN, HIDDEN, bias=False) + + # Block-level RMSNorms (not typically LoRA-targeted but present) + block.adaLN_sa_ln = torch.nn.RMSNorm(HIDDEN) + block.adaLN_mlp_ln = torch.nn.RMSNorm(HIDDEN) + + return block + + +def build_mock_transformer(): + """Build a torch.nn.Module mimicking ``ErnieImageTransformer2DModel``.""" + transformer = _Holder() + transformer.layers = torch.nn.ModuleList([build_ernie_block() for _ in range(N_LAYERS)]) + # Module-level adaLN_modulation: Sequential(SiLU, Linear). + # Real ernie LoRAs target ``adaLN_modulation.1`` (the Linear). + transformer.adaLN_modulation = torch.nn.Sequential( + torch.nn.SiLU(), + torch.nn.Linear(HIDDEN, ADALN_OUT, bias=True), + ) + # final_linear at the model top level - also LoRA-targetable per real fixtures + transformer.final_linear = torch.nn.Linear(HIDDEN, HIDDEN, bias=True) + return transformer + + +class _MockErniePipeline: + """Class name carries 'ErnieImage' so name-based model-type dispatch routes correctly.""" + + def __init__(self, transformer): + self.transformer = transformer + self.text_encoder = None + + +class _MockErnieSdModel: + """Outer wrapper holding pipe + network_layer_mapping.""" + + def __init__(self, pipe): + self.pipe = pipe + self.network_layer_mapping = {} + self.embedding_db = None + self.__class__.__name__ = 'ErnieImagePipeline' + + +def install_mock_pipe(): + """Set shared.sd_model to a mock exposing an ERNIE-Image-shaped transformer. + + Each test re-installs so any prior network_layer_name stamps don't leak. + """ + transformer = build_mock_transformer() + pipe = _MockErniePipeline(transformer) + sd_model = _MockErnieSdModel(pipe) + from modules.modeldata import model_data + model_data.sd_model = sd_model + return sd_model + + +# ============================================================ +# State-dict synthesizers (one per family/format) +# ============================================================ + +RANK_LORA = 8 +LOKR_W1_DIM = 8 + + +def sd_lora_bfl_mlp_gate_proj(): + """BFL LoRA on layers.X.mlp.gate_proj. Mirrors Ernie-Breast-Slider-v1.""" + return { + 'diffusion_model.layers.0.mlp.gate_proj.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.0.mlp.gate_proj.lora_B.weight': torch.randn(FFN_HIDDEN, RANK_LORA), + } + + +def sd_lora_bfl_self_attention(): + """BFL LoRA on the split self_attention.to_q (no fusion in ernie).""" + return { + 'diffusion_model.layers.1.self_attention.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.1.self_attention.to_q.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + 'diffusion_model.layers.1.self_attention.to_q.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_bfl_self_attention_to_out(): + """BFL LoRA on self_attention.to_out.0 (the Linear inside the ModuleList).""" + return { + 'diffusion_model.layers.0.self_attention.to_out.0.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'diffusion_model.layers.0.self_attention.to_out.0.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_bfl_mlp_linear_fc2(): + """BFL LoRA on mlp.linear_fc2 (the down-projection).""" + return { + 'diffusion_model.layers.0.mlp.linear_fc2.lora_A.weight': torch.randn(RANK_LORA, FFN_HIDDEN), + 'diffusion_model.layers.0.mlp.linear_fc2.lora_B.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_kohya_mlp_gate_proj(): + """Kohya flat-underscore LoRA on layers.X.mlp.gate_proj.""" + return { + 'lora_unet_layers_0_mlp_gate_proj.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'lora_unet_layers_0_mlp_gate_proj.lora_up.weight': torch.randn(FFN_HIDDEN, RANK_LORA), + 'lora_unet_layers_0_mlp_gate_proj.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_kohya_adaLN_modulation(): + """Kohya LoRA on module-level adaLN_modulation.1 (the Linear inside Sequential). + + Real fixtures use ``lora_unet_adaLN_modulation_1`` -> the Linear at index 1. + """ + return { + 'lora_unet_adaLN_modulation_1.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'lora_unet_adaLN_modulation_1.lora_up.weight': torch.randn(ADALN_OUT, RANK_LORA), + 'lora_unet_adaLN_modulation_1.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_lora_peft_to_v(): + """PEFT-format LoRA on self_attention.to_v.""" + return { + 'transformer.layers.0.self_attention.to_v.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'transformer.layers.0.self_attention.to_v.lora_up.weight': torch.randn(HIDDEN, RANK_LORA), + } + + +def sd_lora_with_dora_scale(): + """LoRA with dora_scale companion to exercise DoRA threading.""" + return { + 'transformer.layers.0.mlp.up_proj.lora_A.weight': torch.randn(RANK_LORA, HIDDEN), + 'transformer.layers.0.mlp.up_proj.lora_B.weight': torch.randn(FFN_HIDDEN, RANK_LORA), + 'transformer.layers.0.mlp.up_proj.dora_scale': torch.randn(FFN_HIDDEN), + } + + +def sd_lokr_bfl_mlp_gate_proj(): + """BFL LoKR on mlp.gate_proj. Mirrors ERNIE_Anatomy_Male.""" + return { + 'diffusion_model.layers.0.mlp.gate_proj.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM), + 'diffusion_model.layers.0.mlp.gate_proj.lokr_w2': torch.randn(FFN_HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM), + 'diffusion_model.layers.0.mlp.gate_proj.alpha': torch.tensor(float(LOKR_W1_DIM)), + } + + +def sd_lokr_bfl_self_attention(): + """BFL LoKR on self_attention.to_q (no fusion, straight binding).""" + return { + 'diffusion_model.layers.1.self_attention.to_q.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM), + 'diffusion_model.layers.1.self_attention.to_q.lokr_w2': torch.randn(HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM), + 'diffusion_model.layers.1.self_attention.to_q.alpha': torch.tensor(float(LOKR_W1_DIM)), + } + + +def sd_loha_bfl_mlp(): + """LoHA on mlp.linear_fc2.""" + return { + 'diffusion_model.layers.0.mlp.linear_fc2.hada_w1_a': torch.randn(HIDDEN, RANK_LORA), + 'diffusion_model.layers.0.mlp.linear_fc2.hada_w1_b': torch.randn(RANK_LORA, FFN_HIDDEN), + 'diffusion_model.layers.0.mlp.linear_fc2.hada_w2_a': torch.randn(HIDDEN, RANK_LORA), + 'diffusion_model.layers.0.mlp.linear_fc2.hada_w2_b': torch.randn(RANK_LORA, FFN_HIDDEN), + 'diffusion_model.layers.0.mlp.linear_fc2.alpha': torch.tensor(float(RANK_LORA)), + } + + +def sd_oft_lycoris_self_attention(): + """OFT (LyCORIS oft_diag form) on self_attention.to_k.""" + num_blocks = 4 + block_size = HIDDEN // num_blocks + return { + 'diffusion_model.layers.0.self_attention.to_k.oft_blocks': torch.randn(num_blocks, block_size, block_size) * 0.01, + 'diffusion_model.layers.0.self_attention.to_k.oft_diag': torch.ones(num_blocks, block_size), + 'diffusion_model.layers.0.self_attention.to_k.alpha': torch.tensor(0.001), + } + + +# ============================================================ +# Helpers +# ============================================================ + + +class TempLora: + """Context manager: writes a state dict to a temp safetensors file.""" + + def __init__(self, state_dict, name='test'): + self.state_dict = state_dict + self.name = name + self.path = None + + def __enter__(self): + sd = {k: v.contiguous() if isinstance(v, torch.Tensor) else v for k, v in self.state_dict.items()} + fd, self.path = tempfile.mkstemp(suffix='.safetensors', prefix=f'{self.name}_') + os.close(fd) + safetensors.torch.save_file(sd, self.path) + return _MockNetworkOnDisk(self.path, self.name) + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.path and os.path.exists(self.path): + os.unlink(self.path) + + +class _MockNetworkOnDisk: + def __init__(self, filename, name): + self.filename = filename + self.name = name + self.shorthash = '' + self.sd_version = 'unknown' + + +def assert_shape(t: torch.Tensor, expected_shape, label=''): + actual = tuple(t.shape) + assert actual == tuple(expected_shape), f'{label}: shape {actual}, expected {tuple(expected_shape)}' + + +def make_network_for_module(net_module: network.NetworkModule, te_mul: float = 1.0, unet_mul: float = 1.0): + net_module.network.te_multiplier = te_mul + net_module.network.unet_multiplier = unet_mul + return net_module + + +# ============================================================ +# Tests - parsing primitives +# ============================================================ + +CAT_PARSE = category('parse') + + +def test_parse_key_all_prefixes(): + """parse_key recognizes BFL, PEFT, kohya, and bare keys.""" + cases = [ + ('diffusion_model.layers.0.mlp.gate_proj.lora_A.weight', + E.LORA_SUFFIXES, + ('lora_transformer_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')), + ('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 + ('layers.0.mlp.up_proj.lora_A.weight', + E.LORA_SUFFIXES, + ('lora_transformer_layers_0_mlp_up_proj', 'lora_down.weight')), + ('random.unrelated.key', E.LORA_SUFFIXES, None), + ] + for key, suffixes, expected in cases: + got = E.parse_key(key, suffixes) + assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}' + return True + + +def test_marker_disambiguation(): + """Each family's markers reject other families' files.""" + pure_lora = { + 'lora_unet_layers_0_mlp_gate_proj.lora_down.weight': torch.zeros(1, 1), + 'lora_unet_layers_0_mlp_gate_proj.lora_up.weight': torch.zeros(1, 1), + } + assert E.has_marker(pure_lora, E.LORA_MARKERS) + assert not E.has_marker(pure_lora, E.LOKR_MARKERS) + assert not E.has_marker(pure_lora, E.LOHA_MARKERS) + assert not E.has_marker(pure_lora, E.OFT_MARKERS) + + pure_lokr = { + 'diffusion_model.layers.0.mlp.gate_proj.lokr_w1': torch.zeros(1, 1), + 'diffusion_model.layers.0.mlp.gate_proj.lokr_w2': torch.zeros(1, 1), + } + assert E.has_marker(pure_lokr, E.LOKR_MARKERS) + assert not E.has_marker(pure_lokr, E.LORA_MARKERS) + return True + + +# ============================================================ +# Tests - loaders end-to-end +# ============================================================ + +CAT_LOADER = category('loader') + + +def _load_via(try_fn, state_dict, name='test'): + install_mock_pipe() + with TempLora(state_dict, name=name) as nod: + return try_fn(name, nod, lora_scale=1.0) + + +def test_lora_bfl_mlp_gate_proj(): + """BFL LoRA on layers.X.mlp.gate_proj binds straight through.""" + net = _load_via(E.try_load_lora, sd_lora_bfl_mlp_gate_proj()) + assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}' + assert 'lora_transformer_layers_0_mlp_gate_proj' in net.modules + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_lora.NetworkModuleLora) + return True + + +def test_lora_bfl_self_attention(): + """BFL LoRA on self_attention.to_q (no fusion in ernie - straight binding).""" + net = _load_via(E.try_load_lora, sd_lora_bfl_self_attention()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_1_self_attention_to_q' in net.modules + return True + + +def test_lora_bfl_self_attention_to_out(): + """BFL LoRA on self_attention.to_out.0 (Linear inside ModuleList).""" + net = _load_via(E.try_load_lora, sd_lora_bfl_self_attention_to_out()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_0_self_attention_to_out_0' in net.modules + return True + + +def test_lora_bfl_mlp_linear_fc2(): + """BFL LoRA on mlp.linear_fc2 (the down projection).""" + net = _load_via(E.try_load_lora, sd_lora_bfl_mlp_linear_fc2()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_0_mlp_linear_fc2' in net.modules + return True + + +def test_lora_kohya_mlp_gate_proj(): + """Kohya format converges to the same diffusers network_key as BFL.""" + net = _load_via(E.try_load_lora, sd_lora_kohya_mlp_gate_proj()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_0_mlp_gate_proj' in net.modules + return True + + +def test_lora_kohya_adaLN_modulation(): + """Kohya LoRA on the module-level adaLN_modulation.1 (Linear in Sequential).""" + net = _load_via(E.try_load_lora, sd_lora_kohya_adaLN_modulation()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_adaLN_modulation_1' in net.modules + return True + + +def test_lora_peft_to_v(): + """PEFT format binds without rename or chunking.""" + net = _load_via(E.try_load_lora, sd_lora_peft_to_v()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_0_self_attention_to_v' in net.modules + return True + + +def test_lora_dora_threading(): + """dora_scale flows into NetworkModuleLora.dora_scale.""" + net = _load_via(E.try_load_lora, sd_lora_with_dora_scale()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert mod.dora_scale is not None + return True + + +def test_lokr_bfl_mlp_gate_proj(): + """BFL LoKR on mlp.gate_proj binds via NetworkModuleLokr (no chunk class in ernie).""" + net = _load_via(E.try_load_lokr, sd_lokr_bfl_mlp_gate_proj()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_0_mlp_gate_proj' in net.modules + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_lokr.NetworkModuleLokr) + # ernie LoKR never instantiates the chunk variants - no fused targets exist + assert not isinstance(mod, network_lokr.NetworkModuleLokrChunk) + return True + + +def test_lokr_bfl_self_attention(): + """BFL LoKR on self_attention.to_q. No fusion means straight NetworkModuleLokr.""" + net = _load_via(E.try_load_lokr, sd_lokr_bfl_self_attention()) + assert net is not None and len(net.modules) == 1 + assert 'lora_transformer_layers_1_self_attention_to_q' in net.modules + return True + + +def test_loha_bfl_mlp(): + """LoHA on mlp.linear_fc2 binds via NetworkModuleHada (no chunk path).""" + net = _load_via(E.try_load_loha, sd_loha_bfl_mlp()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_hada.NetworkModuleHada) + return True + + +def test_oft_lycoris_no_npe(): + """OFT LyCORIS oft_diag form loads on self_attention.to_k without NoneType errors.""" + net = _load_via(E.try_load_oft, sd_oft_lycoris_self_attention()) + assert net is not None and len(net.modules) == 1 + mod = next(iter(net.modules.values())) + assert isinstance(mod, network_oft.NetworkModuleOFT) + return True + + +# ============================================================ +# Tests - calc_updown shape sanity +# ============================================================ + +CAT_MATH = category('math') + + +def test_lora_calc_updown_shape(): + net = _load_via(E.try_load_lora, sd_lora_bfl_mlp_gate_proj()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(FFN_HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoRA calc_updown') + return True + + +def test_lokr_calc_updown_shape(): + net = _load_via(E.try_load_lokr, sd_lokr_bfl_mlp_gate_proj()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(FFN_HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoKR calc_updown') + return True + + +def test_loha_calc_updown_shape(): + net = _load_via(E.try_load_loha, sd_loha_bfl_mlp()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, FFN_HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='LoHA calc_updown') + return True + + +def test_oft_calc_updown_shape(): + net = _load_via(E.try_load_oft, sd_oft_lycoris_self_attention()) + mod = make_network_for_module(next(iter(net.modules.values()))) + target = torch.randn(HIDDEN, HIDDEN) + updown, _ = mod.calc_updown(target) + assert_shape(updown, target.shape, label='OFT calc_updown') + return True + + +# ============================================================ +# Test runner +# ============================================================ + + +def run_tests(): + t0 = time.time() + + log.warning('=== Parsing primitives ===') + for fn in [test_parse_key_all_prefixes, test_marker_disambiguation]: + run_test(CAT_PARSE, fn) + + log.warning('=== Loaders ===') + for fn in [ + test_lora_bfl_mlp_gate_proj, + test_lora_bfl_self_attention, + test_lora_bfl_self_attention_to_out, + test_lora_bfl_mlp_linear_fc2, + test_lora_kohya_mlp_gate_proj, + test_lora_kohya_adaLN_modulation, + test_lora_peft_to_v, + test_lora_dora_threading, + test_lokr_bfl_mlp_gate_proj, + test_lokr_bfl_self_attention, + test_loha_bfl_mlp, + test_oft_lycoris_no_npe, + ]: + run_test(CAT_LOADER, fn) + + log.warning('=== calc_updown shape sanity ===') + for fn in [ + test_lora_calc_updown_shape, + test_lokr_calc_updown_shape, + test_loha_calc_updown_shape, + test_oft_calc_updown_shape, + ]: + run_test(CAT_MATH, fn) + + elapsed = time.time() - t0 + log.warning('=== Results ===') + total_pass = 0 + total_fail = 0 + for cat, info in results.items(): + status = 'PASS' if info['failed'] == 0 else 'FAIL' + log.info(f' {cat}: {info["passed"]} passed, {info["failed"]} failed [{status}]') + total_pass += info['passed'] + total_fail += info['failed'] + log.warning(f'Total: {total_pass} passed, {total_fail} failed in {elapsed:.2f}s') + return total_fail == 0 + + +if __name__ == '__main__': + ok = run_tests() + sys.exit(0 if ok else 1) From c3b7379fd32989d5b2473f19ecf8a98895e1c013 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 23:14:52 +0100 Subject: [PATCH 10/13] refactor(zimage): migrate to generic native_loader Replaces zimage's four family loaders with thin wrappers binding native_loader's generics to z-image's prefix tuples and resolve_targets. resolve_targets folds the legacy attention.qkv split and attention.out alias rename into the path-resolution step. Fused qkv now emits three ChunkSpec(idx, total=3) entries; attention.out / attention.out.0 / attention.wo aliases collapse to attention.to_out.0. BARE_DIFFUSERS_PREFIXES allows bare paths starting with layers. / noise_refiner. / context_refiner. to pass through to the loader. This matches real Z-Image LoRAs exported via ZImageTransformer2DModel.save_lora_adapter(). LoHA on fused qkv now binds via NetworkModuleHadaChunk (added to shared infra by the flux2 PR) instead of being skipped. Test renamed to test_loha_legacy_fused_qkv_chunked. parse_key returns (prefix_used, base, suffix) instead of the old (network_key, suffix); parse test updated. --- pipelines/z_image/zimage_lora.py | 475 ++++++++-------------------- test/test-zimage-native-adapters.py | 44 ++- 2 files changed, 158 insertions(+), 361 deletions(-) diff --git a/pipelines/z_image/zimage_lora.py b/pipelines/z_image/zimage_lora.py index 2a138859e..49c35f9d8 100644 --- a/pipelines/z_image/zimage_lora.py +++ b/pipelines/z_image/zimage_lora.py @@ -1,379 +1,162 @@ """Z-Image native adapter loader. Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'`` -(``lora_force_diffusers`` off and ``zimage`` 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 ``zimage`` 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` (not yet validated against a real Z-Image adapter) -- OFT via :func:`try_load_oft` (not yet validated against a real Z-Image adapter) - -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``. +Recognized key prefixes: ``diffusion_model.``, ``transformer.``, +``lora_unet_``, or bare paths starting with the known block-level prefixes +(``layers.``, ``noise_refiner.``, ``context_refiner.``). Pre-refactor Z-Image attention layouts (fused ``attention.qkv``, bare -``attention.out`` / ``attention.wo``) are rewritten to the current diffusers -``to_q``/``to_k``/``to_v`` and ``to_out.0``. For LoRA the fused qkv up-weight -is chunked along dim 0 at load time. For LoKR the split is deferred to apply -time via :class:`NetworkModuleLokrChunk`, which materializes ``kron(w1, w2)`` -once per forward pass and returns the designated slice. +``attention.out`` / ``attention.wo``) are rewritten by :func:`resolve_targets` +to the current diffusers ``to_q``/``to_k``/``to_v`` and ``to_out.0``. For +LoRA the fused qkv up-weight is chunked at load time; for LoKR the split is +deferred to apply time via :class:`network_lokr.NetworkModuleLokrChunk`. Fused ``attention.qkv`` for LoHA and OFT is skipped with a warning: no -``NetworkModuleHadaChunk`` exists, and an OFT block structure is tied to the -target module's ``out_features``, so a Q/K/V split is not a drop-in. +slice variant exists for LoHA's Hadamard product on Linear targets, and an +OFT block structure is tied to the target module's ``out_features`` so a +Q/K/V split is not a drop-in. """ -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 +from modules.lora.native_loader import ChunkSpec -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 -# Presence of any of these substrings anywhere in a key marks a file as belonging to that family. -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", -} - -ATTENTION_OUT_ALIASES = ("attention_out", "attention_out_0", "attention_wo") -ATTENTION_OUT_TARGET = "attention_to_out_0" -ATTENTION_QKV_SUFFIX = "attention_qkv" -ATTENTION_QKV_TARGETS = ("attention_to_q", "attention_to_k", "attention_to_v") +BARE_DIFFUSERS_PREFIXES = ("layers.", "noise_refiner.", "context_refiner.") -def try_load_lora(name, network_on_disk, lora_scale): - """Try loading a Z-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 +# === Re-exports for test/back-compat === - mapping = resolve_mapping() - net = new_network(name, network_on_disk) +LORA_SUFFIXES = native_loader.LORA_SUFFIXES +LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES +LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES +OFT_SUFFIXES = native_loader.OFT_SUFFIXES - groups = group_by_suffixes(state_dict, LORA_SUFFIXES) - groups = expand_legacy_attention_lora(groups) +LORA_MARKERS = native_loader.LORA_MARKERS +LOKR_MARKERS = native_loader.LOKR_MARKERS +LOHA_MARKERS = native_loader.LOHA_MARKERS +OFT_MARKERS = native_loader.OFT_MARKERS - 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 a Z-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) - groups, chunk_info = expand_legacy_attention_lokr(groups) - - 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) - chunked = chunk_info.get(network_key) - if chunked is not None: - idx, num = chunked - net.modules[network_key] = network_lokr.NetworkModuleLokrChunk(net, nw, idx, num) - else: - 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 a Z-Image LoHA as native modules. Fused attention.qkv groups are skipped.""" - 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) - groups = rename_attention_out(groups) - - unmapped = 0 - skipped_qkv = 0 - for network_key, w in groups.items(): - if network_key.endswith("_" + ATTENTION_QKV_SUFFIX): - log.warning(f'Network load: type=LoHA name="{name}" key={network_key} fused qkv skipped (unsupported)') - skipped_qkv += 1 - continue - 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, skipped=skipped_qkv) - - -def try_load_oft(name, network_on_disk, lora_scale): - """Try loading a Z-Image OFT adapter as native modules. Fused attention.qkv groups are skipped.""" - 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) - groups = rename_attention_out(groups) - - unmapped = 0 - skipped_qkv = 0 - for network_key, w in groups.items(): - if network_key.endswith("_" + ATTENTION_QKV_SUFFIX): - log.warning(f'Network load: type=OFT name="{name}" key={network_key} fused qkv skipped (unsupported)') - skipped_qkv += 1 - continue - 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, skipped=skipped_qkv) - - -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, skipped=0): - if len(net.modules) == 0: - if unmapped or mismatch or skipped: - log.debug( - f'Network load: type={family} name="{name}" native no-match' - f' unmapped={unmapped} mismatch={mismatch} skipped={skipped}' - ) - return None - log.debug( - f'Network load: type={family} name="{name}" native modules={len(net.modules)}' - f' unmapped={unmapped} mismatch={mismatch} skipped={skipped} 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 - - 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 - - suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix) - network_key = 'lora_transformer_' + base.replace('.', '_') - return network_key, suffix + """Z-Image-bound :func:`native_loader.parse_key`.""" + return native_loader.parse_key( + key, suffixes, + prefixes=KNOWN_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + ) -def rename_attention_out(groups): - """Rename legacy ``attention.out`` / ``attention.wo`` keys to ``attention.to_out.0``.""" - out: dict[str, dict[str, torch.Tensor]] = {} - for key, w in groups.items(): - new_key = None - for alias in ATTENTION_OUT_ALIASES: - if key.endswith("_" + alias): - new_key = key[: -len(alias)] + ATTENTION_OUT_TARGET - break - out[new_key or key] = w - return out +def group_by_suffixes(state_dict, suffixes): + """Z-Image-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, + ) -def expand_legacy_attention_lora(groups): - """Rename attention.out, then split fused attention.qkv into three per-projection groups. +# === Target resolution (arch-specific) === - The down-weight is shared across Q/K/V and the up-weight is chunked along - dim 0 (the concatenated output dim in the fused layout). + +def resolve_targets(prefix_used, base): + """Return ``[(diffusers_path, ChunkSpec | None), ...]`` for a parsed group key. + + Handles two legacy Z-Image attention layouts: + + - Fused ``attention.qkv`` is split into three Q/K/V targets, each carrying + a :class:`ChunkSpec` for the loader to apply. + - ``attention.out`` / ``attention.out.0`` / ``attention.wo`` are aliased to + the current diffusers ``attention.to_out.0`` path. + + Everything else (modern split-attention paths, MLP, norms, embedders) is + returned verbatim. """ - groups = rename_attention_out(groups) - out: dict[str, dict[str, torch.Tensor]] = {} - for key, w in groups.items(): - if not key.endswith("_" + ATTENTION_QKV_SUFFIX): - out[key] = w - continue - stem = key[: -len(ATTENTION_QKV_SUFFIX)] - down = w.get("lora_down.weight") - up = w.get("lora_up.weight") - if down is None or up is None or up.shape[0] % 3 != 0: - out[key] = w - continue - chunks = torch.chunk(up, 3, dim=0) - alpha = w.get("alpha") - dora = w.get("dora_scale") - for target, chunk in zip(ATTENTION_QKV_TARGETS, chunks): - split_key = stem + target - split = { - "lora_down.weight": down, - "lora_up.weight": chunk.contiguous(), - } - if alpha is not None: - split["alpha"] = alpha - if dora is not None: - split["dora_scale"] = dora - out[split_key] = split - return out + if prefix_used == "lora_unet_": + return _underscore_to_diffusers_targets(base) + if prefix_used == "transformer.": + return [(base, None)] + if prefix_used == BARE_DIFFUSERS_PREFIX_USED: + return [(base, None)] + if prefix_used in (None, "diffusion_model."): + return _dotted_to_diffusers_targets(base) + return [] -def expand_legacy_attention_lokr(groups): - """Rename attention.out, then split fused attention.qkv for LoKR. +def _dotted_to_diffusers_targets(base): + """For BFL / bare-BFL keys like ``layers.0.attention.qkv``.""" + if base.endswith(".attention.qkv"): + stem = base[:-len(".attention.qkv")] + return [ + (f"{stem}.attention.to_q", ChunkSpec(idx=0, total=3)), + (f"{stem}.attention.to_k", ChunkSpec(idx=1, total=3)), + (f"{stem}.attention.to_v", ChunkSpec(idx=2, total=3)), + ] + for alias in (".attention.out.0", ".attention.out", ".attention.wo"): + if base.endswith(alias): + stem = base[:-len(alias)] + return [(f"{stem}.attention.to_out.0", None)] + return [(base, None)] - For LoKR the split is deferred to apply time via ``NetworkModuleLokrChunk`` - (returned via a parallel ``chunk_info`` dict keyed by the split network key). - The three resulting groups share the same tensor dict by shallow copy; the - chunk module computes ``kron(w1, w2)`` once and slices the designated row - range. - Returns ``(groups, chunk_info)`` where ``chunk_info[network_key] == (index, num)``. - """ - groups = rename_attention_out(groups) - out: dict[str, dict[str, torch.Tensor]] = {} - chunk_info: dict[str, tuple[int, int]] = {} - for key, w in groups.items(): - if not key.endswith("_" + ATTENTION_QKV_SUFFIX): - out[key] = w - continue - stem = key[: -len(ATTENTION_QKV_SUFFIX)] - for i, target in enumerate(ATTENTION_QKV_TARGETS): - split_key = stem + target - out[split_key] = dict(w) - chunk_info[split_key] = (i, 3) - return out, chunk_info +def _underscore_to_diffusers_targets(base): + """For kohya flat-underscore keys like ``layers_0_attention_qkv``.""" + if base.endswith("_attention_qkv"): + stem = base[:-len("_attention_qkv")] + return [ + (f"{stem}_attention_to_q", ChunkSpec(idx=0, total=3)), + (f"{stem}_attention_to_k", ChunkSpec(idx=1, total=3)), + (f"{stem}_attention_to_v", ChunkSpec(idx=2, total=3)), + ] + for alias in ("_attention_out_0", "_attention_out", "_attention_wo"): + if base.endswith(alias): + stem = base[:-len(alias)] + return [(f"{stem}_attention_to_out_0", None)] + return [(base, None)] + + +# === 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="zimage", +) + + +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 Z-Image 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-zimage-native-adapters.py b/test/test-zimage-native-adapters.py index b75a0ab00..79c838b90 100644 --- a/test/test-zimage-native-adapters.py +++ b/test/test-zimage-native-adapters.py @@ -414,25 +414,27 @@ CAT_PARSE = category('parse') def test_parse_key_all_prefixes(): - """parse_key recognizes BFL, PEFT, kohya, and bare keys with the right base+suffix.""" + """parse_key recognizes BFL, PEFT, kohya, and bare-diffusers keys. + + Returns (prefix_used, base, suffix) - prefix_used is the matched + KNOWN_PREFIXES element, BARE_DIFFUSERS_PREFIX_USED for bare paths + matching BARE_DIFFUSERS_PREFIXES, or None when no prefix is recognized. + """ + bd = Z.BARE_DIFFUSERS_PREFIX_USED cases = [ - # BFL prefix -> dotted base, lora_A normalized to lora_down ('diffusion_model.layers.0.attention.to_q.lora_A.weight', Z.LORA_SUFFIXES, - ('lora_transformer_layers_0_attention_to_q', 'lora_down.weight')), - # PEFT prefix + ('diffusion_model.', 'layers.0.attention.to_q', 'lora_down.weight')), ('transformer.layers.0.attention.to_v.lora_B.weight', Z.LORA_SUFFIXES, - ('lora_transformer_layers_0_attention_to_v', 'lora_up.weight')), - # kohya flat underscore form + ('transformer.', 'layers.0.attention.to_v', 'lora_up.weight')), ('lora_unet_layers_0_attention_to_k.lora_down.weight', Z.LORA_SUFFIXES, - ('lora_transformer_layers_0_attention_to_k', 'lora_down.weight')), - # Bare dotted (no recognized prefix) + ('lora_unet_', 'layers_0_attention_to_k', 'lora_down.weight')), + # Bare path starting with a known block prefix ('layers.0.attention.to_out.0.lora_A.weight', Z.LORA_SUFFIXES, - ('lora_transformer_layers_0_attention_to_out_0', 'lora_down.weight')), - # Unrelated keys reject cleanly + (bd, 'layers.0.attention.to_out.0', 'lora_down.weight')), ('random.unrelated.key', Z.LORA_SUFFIXES, None), ] for key, suffixes, expected in cases: @@ -580,11 +582,23 @@ def test_loha_bfl_proj(): return True -def test_loha_legacy_fused_qkv_skipped(): - """LoHA on legacy fused attention.qkv is skipped with a warning (no chunk variant).""" +def test_loha_legacy_fused_qkv_chunked(): + """LoHA on legacy fused attention.qkv emits 3 HadaChunk modules. + + Post-migration the generic LoHA loader uses NetworkModuleHadaChunk for + equal-chunks dispatch (the chunk class was added in the flux2 PR and is + now shared infrastructure). + """ net = _load_via(Z.try_load_loha, sd_loha_legacy_fused_qkv_skipped()) - # Either no net returned (nothing matched) or net with zero modules - assert net is None or len(net.modules) == 0, f'expected no modules, got {net.modules if net else None}' + assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}' + expected = { + 'lora_transformer_layers_0_attention_to_q', + 'lora_transformer_layers_0_attention_to_k', + 'lora_transformer_layers_0_attention_to_v', + } + assert set(net.modules) == expected, f'got {set(net.modules)}' + for mod in net.modules.values(): + assert isinstance(mod, network_hada.NetworkModuleHadaChunk) return True @@ -689,7 +703,7 @@ def run_tests(): test_lokr_bfl_adaln, test_lokr_legacy_fused_qkv_chunked, test_loha_bfl_proj, - test_loha_legacy_fused_qkv_skipped, + test_loha_legacy_fused_qkv_chunked, test_oft_lycoris_no_npe, test_oft_legacy_fused_qkv_skipped, ]: From e45c23c0319c0cc1f1abaed83387e0c5d66ead98 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 23:31:17 +0100 Subject: [PATCH 11/13] refactor(chroma): migrate to generic native_loader Replaces chroma's four family loaders with thin wrappers binding native_loader's generics to chroma's prefix tuples and resolve_targets. resolve_targets folds the Flux-to-diffusers rename table and the two fused-weight splits into one path-resolution step: - img_attn.qkv / txt_attn.qkv: ChunkSpec(idx, total=3) for equal Q/K/V - single_blocks.linear1: ChunkSpec(start, end) for the unequal Q/K/V/proj_mlp partition QKV_DIMS and LINEAR1_DIMS stay as module-level constants (tests patch them for the scaled-down mock). Behavior changes: - LoKR on equal-chunk QKV now dispatches to NetworkModuleLokrChunk instead of NetworkModuleLokrSliceChunk (the slice variant was used pre-migration for both forms since chroma had no equal-chunk path). - LoHA on fused img_attn.qkv now binds via NetworkModuleHadaChunk instead of being skipped; the shared HadaChunk added in the flux2 PR is general for equal-chunks dispatch. parse_key returns (prefix_used, base, suffix) instead of the old (flat_key, suffix); rename happens in resolve_targets. test_static_rename_table replaced with test_resolve_targets_static_renames driving the same remappings through the new interface. --- pipelines/chroma/chroma_lora.py | 640 +++++++++------------------- test/test-chroma-native-adapters.py | 96 +++-- 2 files changed, 241 insertions(+), 495 deletions(-) diff --git a/pipelines/chroma/chroma_lora.py b/pipelines/chroma/chroma_lora.py index c92635948..8354a98b5 100644 --- a/pipelines/chroma/chroma_lora.py +++ b/pipelines/chroma/chroma_lora.py @@ -1,498 +1,238 @@ """Chroma native adapter loader. Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'`` -(``lora_force_diffusers`` off and ``chroma`` 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 ``chroma`` 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` (fused groups skipped, no chunk variant) -- OFT via :func:`try_load_oft` (fused groups skipped, no chunk variant) +Recognized key prefixes: ``diffusion_model.``, ``transformer.``, +``lora_unet_``, plus bare BFL paths (``double_blocks.`` / ``single_blocks.``) +and bare diffusers paths (``transformer_blocks.`` / +``single_transformer_blocks.`` / ``distilled_guidance_layer.``). -Recognized key prefixes: ``diffusion_model.``, ``transformer.``, ``lora_unet_``. -Diffusers-PEFT ``lora_A``/``lora_B`` are normalized to ``lora_down``/``lora_up``. +Chroma LoRAs are trained against the Flux block layout regardless of save +format. The diffusers ``ChromaTransformer2DModel`` exposes split-attention +modules; :func:`resolve_targets` rewrites Flux-layout keys to diffusers +paths and emits ``ChunkSpec`` entries for fused targets. -Chroma LoRAs are trained against the Flux block layout regardless of which -key style they save in: +Fused weight handling: -- ``double_blocks.{i}.{img,txt}_attn.{proj,qkv}`` -- ``double_blocks.{i}.{img,txt}_mlp.{0,2}`` -- ``single_blocks.{i}.{linear1,linear2}`` - -The diffusers ``ChromaTransformer2DModel`` exposes split-attention modules at: - -- ``transformer_blocks.{i}.attn.{to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out}`` -- ``transformer_blocks.{i}.{ff,ff_context}.net.{0.proj,2}`` -- ``single_transformer_blocks.{i}.attn.{to_q,to_k,to_v}`` -- ``single_transformer_blocks.{i}.{proj_mlp,proj_out}`` - -This loader path-rewrites Flux-layout keys to diffusers names and expands -fused QKV at load time (LoRA) or apply time (LoKR via -:class:`NetworkModuleLokrSliceChunk`). For LoHA/OFT the fused groups are -skipped with a warning. +- ``double_blocks.{i}.img_attn.qkv`` / ``txt_attn.qkv``: three equal Q/K/V + chunks (``ChunkSpec(idx, total=3)``). +- ``single_blocks.{i}.linear1``: four chunks Q/K/V/proj_mlp at unequal dims + ``[HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN]`` (``ChunkSpec(start, end)``). Chroma's modulation generator is the central ``distilled_guidance_layer`` approximator (replacing Flux's per-block ``norm1.linear``). The pruned -AdaLN classes have no ``.linear`` submodule, so any ``_mod_lin`` / -``_modulation_lin`` keys naturally land in ``unmapped`` and are reported. -LoRAs that target the approximator itself work without special casing, -since ``distilled_guidance_layer.<...>`` is a real module path that -``assign_network_names_to_compvis_modules`` registers. +``ChromaAdaLayerNormZeroPruned`` classes have no ``.linear`` submodule, so +any ``_mod_lin`` / ``_modulation_lin`` keys land in ``unmapped``. LoRAs +targeting the approximator pass through unchanged. """ -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 +from modules.lora.native_loader import ChunkSpec -KNOWN_PREFIXES = ("diffusion_model.", "transformer.", "lora_unet_") +# === Arch-specific prefix configuration === -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_FLUX_PREFIXES = ("double_blocks.", "single_blocks.") + +BARE_DIFFUSERS_PREFIXES = ( + "transformer_blocks.", "single_transformer_blocks.", + "distilled_guidance_layer.", ) -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", -} +# === Fused weight dims === +# Defaults match Chroma1-HD (``inner_dim = num_attention_heads * +# attention_head_dim = 24 * 128 = 3072``; ``mlp_hidden = 12288``). Tests +# patch these to the mock's scale via direct module-level assignment. -# Default block counts for Chroma1-HD; overridden at runtime from the live -# transformer's config when available. -DEFAULT_NUM_DOUBLE_LAYERS = 19 -DEFAULT_NUM_SINGLE_LAYERS = 38 - -# Fused QKV split dims. Single blocks fuse Q/K/V plus proj_mlp into linear1; -# the last chunk is unequal (12288 vs 3072 for Q/K/V). QKV_DIMS = [3072, 3072, 3072] LINEAR1_DIMS = [3072, 3072, 3072, 12288] -# Static (non-fused) renames from underscore-flat Flux-layout paths to -# underscore-flat diffusers paths. Built per-call by format()-ing the layer index. -DOUBLE_RENAME_TEMPLATES = { - 'double_blocks_{i}_img_attn_proj': 'transformer_blocks_{i}_attn_to_out_0', - 'double_blocks_{i}_img_mlp_0': 'transformer_blocks_{i}_ff_net_0_proj', - 'double_blocks_{i}_img_mlp_2': 'transformer_blocks_{i}_ff_net_2', - 'double_blocks_{i}_txt_attn_proj': 'transformer_blocks_{i}_attn_to_add_out', - 'double_blocks_{i}_txt_mlp_0': 'transformer_blocks_{i}_ff_context_net_0_proj', - 'double_blocks_{i}_txt_mlp_2': 'transformer_blocks_{i}_ff_context_net_2', -} -SINGLE_RENAME_TEMPLATES = { - 'single_blocks_{i}_linear2': 'single_transformer_blocks_{i}_proj_out', -} -# Fused-target qkv mappings. Double-block qkv fans out to img-side (to_*) and -# context-side (add_*_proj). Single-block linear1 fans out to single attn and -# proj_mlp. -DOUBLE_IMG_QKV_TARGETS = ('attn_to_q', 'attn_to_k', 'attn_to_v') -DOUBLE_TXT_QKV_TARGETS = ('attn_add_q_proj', 'attn_add_k_proj', 'attn_add_v_proj') -SINGLE_LINEAR1_TARGETS = ('attn_to_q', 'attn_to_k', 'attn_to_v', 'proj_mlp') +# === 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 build_static_rename(num_double, num_single): - """Return {flux_flat_name: diffusers_flat_name} for non-fused paths.""" - out = {} - for i in range(num_double): - for src, dst in DOUBLE_RENAME_TEMPLATES.items(): - out[src.format(i=i)] = dst.format(i=i) - for i in range(num_single): - for src, dst in SINGLE_RENAME_TEMPLATES.items(): - out[src.format(i=i)] = dst.format(i=i) - return out +LORA_MARKERS = native_loader.LORA_MARKERS +LOKR_MARKERS = native_loader.LOKR_MARKERS +LOHA_MARKERS = native_loader.LOHA_MARKERS +OFT_MARKERS = native_loader.OFT_MARKERS - -def get_block_counts(): - """Read num_layers / num_single_layers from the live transformer, with fallback.""" - sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) - transformer = getattr(sd_model, 'transformer', None) - config = getattr(transformer, 'config', None) - num_double = getattr(config, 'num_layers', DEFAULT_NUM_DOUBLE_LAYERS) if config is not None else DEFAULT_NUM_DOUBLE_LAYERS - num_single = getattr(config, 'num_single_layers', DEFAULT_NUM_SINGLE_LAYERS) if config is not None else DEFAULT_NUM_SINGLE_LAYERS - return num_double, num_single - - -def try_load_lora(name, network_on_disk, lora_scale): - """Try loading a Chroma 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 - - mapping = resolve_mapping() - net = new_network(name, network_on_disk) - static_rename = build_static_rename(*get_block_counts()) - - groups = group_by_suffixes(state_dict, LORA_SUFFIXES) - groups = expand_chroma_fused_lora(groups) - groups = apply_static_rename(groups, static_rename) - - 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 a Chroma 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) - static_rename = build_static_rename(*get_block_counts()) - - groups = group_by_suffixes(state_dict, LOKR_SUFFIXES) - groups, slice_info = expand_chroma_fused_lokr(groups) - groups = apply_static_rename(groups, static_rename) - # Mirror apply_static_rename: the slice_info dict must use the same final - # network keys as ``groups`` (static rename applied + ``lora_transformer_`` - # prefix) so the per-target lookup in the loop matches. - slice_info = {'lora_transformer_' + static_rename.get(k, k): v for k, v in slice_info.items()} - - 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) - rng = slice_info.get(network_key) - if rng is not None: - start, end = rng - net.modules[network_key] = network_lokr.NetworkModuleLokrSliceChunk(net, nw, start, end) - else: - 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 a Chroma LoHA as native modules. Fused qkv/linear1 groups are skipped.""" - 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) - static_rename = build_static_rename(*get_block_counts()) - - groups = group_by_suffixes(state_dict, LOHA_SUFFIXES) - groups, skipped = drop_chroma_fused_groups(groups, family='LoHA', name=name) - groups = apply_static_rename(groups, static_rename) - - 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, skipped=skipped) - - -def try_load_oft(name, network_on_disk, lora_scale): - """Try loading a Chroma OFT adapter as native modules. Fused qkv/linear1 groups are skipped.""" - 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) - static_rename = build_static_rename(*get_block_counts()) - - groups = group_by_suffixes(state_dict, OFT_SUFFIXES) - groups, skipped = drop_chroma_fused_groups(groups, family='OFT', name=name) - groups = apply_static_rename(groups, static_rename) - - 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, skipped=skipped) - - -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, skipped=0): - if len(net.modules) == 0: - if unmapped or mismatch or skipped: - log.debug( - f'Network load: type={family} name="{name}" native no-match' - f' unmapped={unmapped} mismatch={mismatch} skipped={skipped}' - ) - return None - log.debug( - f'Network load: type={family} name="{name}" native modules={len(net.modules)}' - f' unmapped={unmapped} mismatch={mismatch} skipped={skipped} 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 ``{flat_key: {suffix: tensor, ...}}`` where the flat key follows - the Flux pre-rename layout (``double_blocks_{i}_img_attn_qkv`` etc.). - """ - groups: dict[str, dict[str, torch.Tensor]] = {} - for key, value in state_dict.items(): - parsed = parse_key(key, suffixes) - if parsed is None: - continue - flat_key, suffix = parsed - slot = groups.get(flat_key) - if slot is None: - slot = {} - groups[flat_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): - """Strip prefix and suffix, return (flat_key, normalized_suffix) or None.""" - stripped = key - for p in KNOWN_PREFIXES: - if key.startswith(p): - stripped = key[len(p):] - break - - 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 - - suffix = SUFFIX_NORMALIZE.get(matched_suffix, matched_suffix) - flat_key = base.replace('.', '_') - return flat_key, suffix + """Chroma-bound :func:`native_loader.parse_key`.""" + return native_loader.parse_key( + key, suffixes, + prefixes=KNOWN_PREFIXES, + bare_prefixes=BARE_FLUX_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + ) -def expand_chroma_fused_lora(groups): - """Split fused QKV / linear1 LoRA groups into their per-projection components. +def group_by_suffixes(state_dict, suffixes): + """Chroma-bound :func:`native_loader.group_by_suffixes`.""" + return native_loader.group_by_suffixes( + state_dict, suffixes, + prefixes=KNOWN_PREFIXES, + bare_prefixes=BARE_FLUX_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + ) - Chroma LoRAs are trained against Flux's fused-attention layout, while the - diffusers ``ChromaTransformer2DModel`` exposes split modules. The - fused-attention LoRA convention shares ``down`` across the splits and - concatenates ``up`` along dim 0; the inverse here splits ``up`` by the - per-target dim list while copying ``down`` to each child. + +# === Target resolution (arch-specific) === + + +def resolve_targets(prefix_used, base): + """Rewrite Flux-layout keys to diffusers paths; emit ChunkSpec on fused targets. + + Dispatches on ``prefix_used``: + + - ``lora_unet_``: kohya underscore-flat Flux path; parse block type/index + and module suffix, rename to diffusers. + - ``diffusion_model.`` or bare BFL (None): dotted Flux path; same rewrite. + - ``transformer.`` or bare-diffusers: already a diffusers path; passthrough. """ - out: dict[str, dict[str, torch.Tensor]] = {} - for key, w in groups.items(): - if key.endswith('_img_attn_qkv') or key.endswith('_txt_attn_qkv'): - stem = key[:-len('img_attn_qkv')] if key.endswith('_img_attn_qkv') else key[:-len('txt_attn_qkv')] - targets = DOUBLE_IMG_QKV_TARGETS if key.endswith('_img_attn_qkv') else DOUBLE_TXT_QKV_TARGETS - block_prefix = stem.replace('double_blocks_', 'transformer_blocks_') - split_groups = split_fused_lora_group(w, QKV_DIMS, [block_prefix + t for t in targets]) - if split_groups is None: - out[key] = w - continue - out.update(split_groups) - elif key.endswith('_linear1') and 'single_blocks' in key: - stem = key[:-len('linear1')] - block_prefix = stem.replace('single_blocks_', 'single_transformer_blocks_') - split_groups = split_fused_lora_group(w, LINEAR1_DIMS, [block_prefix + t for t in SINGLE_LINEAR1_TARGETS]) - if split_groups is None: - out[key] = w - continue - out.update(split_groups) - else: - out[key] = w - return out + if prefix_used == "transformer.": + return [(base, None)] + if prefix_used == BARE_DIFFUSERS_PREFIX_USED: + return [(base, None)] + if prefix_used == "lora_unet_": + return _kohya_to_diffusers(base) + if prefix_used in (None, "diffusion_model."): + return _bfl_to_diffusers(base) + return [] -def split_fused_lora_group(w, dims, target_keys): - """Split a fused LoRA (down, up) into per-target groups by row dim. +def _kohya_to_diffusers(base): + """For kohya keys like ``double_blocks_0_img_attn_qkv`` or ``single_blocks_5_linear1``.""" + if base.startswith("double_blocks_"): + rest = base[len("double_blocks_"):] + idx, _, suffix = rest.partition("_") + return _double_block_targets(idx, suffix) + if base.startswith("single_blocks_"): + rest = base[len("single_blocks_"):] + idx, _, suffix = rest.partition("_") + return _single_block_targets(idx, suffix) + return [(base, None)] - Returns ``{target_key: {suffix: tensor, ...}}`` or ``None`` if the input is - malformed (missing tensors, up-weight rows don't sum to dims). - """ - down = w.get('lora_down.weight') - up = w.get('lora_up.weight') - if down is None or up is None: - return None - if up.shape[0] != sum(dims): - return None - alpha = w.get('alpha') - dora = w.get('dora_scale') - bias = w.get('bias') - scale = w.get('scale') - out: dict[str, dict[str, torch.Tensor]] = {} + +def _bfl_to_diffusers(base): + """For BFL dotted keys like ``double_blocks.0.img_attn.qkv``.""" + parts = base.split(".") + if len(parts) < 3: + return [(base, None)] + block_type, block_idx, module_suffix = parts[0], parts[1], ".".join(parts[2:]) + # Normalize the dotted module suffix to the underscore form the dispatch tables use. + suffix_key = module_suffix.replace(".", "_") + if block_type == "double_blocks": + return _double_block_targets(block_idx, suffix_key) + if block_type == "single_blocks": + return _single_block_targets(block_idx, suffix_key) + return [(base, None)] + + +# Static non-fused renames (underscore-keyed for dispatch from either kohya or BFL paths). +_DOUBLE_STATIC = { + "img_attn_proj": "attn.to_out.0", + "txt_attn_proj": "attn.to_add_out", + "img_mlp_0": "ff.net.0.proj", + "img_mlp_2": "ff.net.2", + "txt_mlp_0": "ff_context.net.0.proj", + "txt_mlp_2": "ff_context.net.2", +} +_SINGLE_STATIC = { + "linear2": "proj_out", +} + + +def _double_block_targets(block_idx, suffix_key): + """Resolve a double_blocks.{i}. target to diffusers paths.""" + if suffix_key == "img_attn_qkv": + return _split_double_qkv(block_idx, ("to_q", "to_k", "to_v")) + if suffix_key == "txt_attn_qkv": + return _split_double_qkv(block_idx, ("add_q_proj", "add_k_proj", "add_v_proj")) + if suffix_key in _DOUBLE_STATIC: + return [(f"transformer_blocks.{block_idx}.{_DOUBLE_STATIC[suffix_key]}", None)] + return [] + + +def _single_block_targets(block_idx, suffix_key): + """Resolve a single_blocks.{i}. target to diffusers paths.""" + if suffix_key == "linear1": + return _split_single_linear1(block_idx) + if suffix_key in _SINGLE_STATIC: + return [(f"single_transformer_blocks.{block_idx}.{_SINGLE_STATIC[suffix_key]}", None)] + return [] + + +def _split_double_qkv(block_idx, attn_keys): + """Three equal-chunk ChunkSpec entries for fused img_attn.qkv / txt_attn.qkv.""" + return [ + (f"transformer_blocks.{block_idx}.attn.{k}", ChunkSpec(idx=i, total=len(attn_keys))) + for i, k in enumerate(attn_keys) + ] + + +def _split_single_linear1(block_idx): + """Four unequal-range ChunkSpec entries for fused linear1 (Q/K/V/proj_mlp).""" + target_keys = ( + f"single_transformer_blocks.{block_idx}.attn.to_q", + f"single_transformer_blocks.{block_idx}.attn.to_k", + f"single_transformer_blocks.{block_idx}.attn.to_v", + f"single_transformer_blocks.{block_idx}.proj_mlp", + ) + targets = [] start = 0 - for tk, d in zip(target_keys, dims): - chunk_up = up[start:start + d].contiguous() + for d, target in zip(LINEAR1_DIMS, target_keys): + targets.append((target, ChunkSpec(start=start, end=start + d))) start += d - slot = {'lora_down.weight': down, 'lora_up.weight': chunk_up} - if alpha is not None: - slot['alpha'] = alpha - if dora is not None: - slot['dora_scale'] = dora - if bias is not None: - slot['bias'] = bias - if scale is not None: - slot['scale'] = scale - out[tk] = slot - return out + return targets -def expand_chroma_fused_lokr(groups): - """Mark fused QKV / linear1 LoKR groups as slice-chunked. - - LoKR factorizations don't compose with row-splitting at load time without - materializing the full Kronecker product. Instead, each target gets a - shallow copy of the same tensor dict, plus an entry in ``slice_info`` that - drives :class:`NetworkModuleLokrSliceChunk` to slice rows lazily on each - forward pass. - """ - out: dict[str, dict[str, torch.Tensor]] = {} - slice_info: dict[str, tuple[int, int]] = {} - for key, w in groups.items(): - if key.endswith('_img_attn_qkv') or key.endswith('_txt_attn_qkv'): - stem = key[:-len('img_attn_qkv')] if key.endswith('_img_attn_qkv') else key[:-len('txt_attn_qkv')] - targets = DOUBLE_IMG_QKV_TARGETS if key.endswith('_img_attn_qkv') else DOUBLE_TXT_QKV_TARGETS - block_prefix = stem.replace('double_blocks_', 'transformer_blocks_') - assign_lokr_slices(out, slice_info, w, QKV_DIMS, [block_prefix + t for t in targets]) - elif key.endswith('_linear1') and 'single_blocks' in key: - stem = key[:-len('linear1')] - block_prefix = stem.replace('single_blocks_', 'single_transformer_blocks_') - assign_lokr_slices(out, slice_info, w, LINEAR1_DIMS, [block_prefix + t for t in SINGLE_LINEAR1_TARGETS]) - else: - out[key] = w - return out, slice_info +# === Native loaders (thin wrappers over native_loader generics) === -def assign_lokr_slices(out, slice_info, w, dims, target_keys): - start = 0 - for tk, d in zip(target_keys, dims): - out[tk] = dict(w) - slice_info[tk] = (start, start + d) - start += d +_BIND_KWARGS = dict( + resolve_targets=resolve_targets, + prefixes=KNOWN_PREFIXES, + bare_prefixes=BARE_FLUX_PREFIXES, + bare_diffusers_prefixes=BARE_DIFFUSERS_PREFIXES, + arch_name="chroma", +) -def drop_chroma_fused_groups(groups, family, name): - """Remove fused QKV / linear1 groups (no chunk variant for LoHA/OFT).""" - out: dict[str, dict[str, torch.Tensor]] = {} - skipped = 0 - for key, w in groups.items(): - is_fused_qkv = key.endswith('_img_attn_qkv') or key.endswith('_txt_attn_qkv') - is_fused_linear1 = key.endswith('_linear1') and 'single_blocks' in key - if is_fused_qkv or is_fused_linear1: - log.warning(f'Network load: type={family} name="{name}" key={key} fused group skipped (unsupported)') - skipped += 1 - continue - out[key] = w - return out, skipped +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 apply_static_rename(groups, static_rename): - """Rewrite Flux-layout flat keys to diffusers flat keys, then prepend ``lora_transformer_``. +def try_load_lokr(name, network_on_disk, lora_scale): + return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS) - Keys without an entry in ``static_rename`` are passed through unchanged - (they may already be diffusers paths from PEFT-style files, or they may - target the ``distilled_guidance_layer`` approximator). The final - ``lora_transformer_`` prefix is added uniformly to match the format - ``assign_network_names_to_compvis_modules`` registers. - """ - out: dict[str, dict[str, torch.Tensor]] = {} - for key, w in groups.items(): - renamed = static_rename.get(key, key) - out['lora_transformer_' + renamed] = w - return out + +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 Chroma 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-chroma-native-adapters.py b/test/test-chroma-native-adapters.py index 087b8db82..e21e51ab5 100644 --- a/test/test-chroma-native-adapters.py +++ b/test/test-chroma-native-adapters.py @@ -528,25 +528,22 @@ CAT_PARSE = category('parse') def test_parse_key_all_prefixes(): - """parse_key recognizes BFL, PEFT, kohya keys with the right flat_key + suffix. - - Chroma's parse_key returns (flat_key, suffix) where flat_key is the - underscored Flux-layout path (before static rename). - """ + """parse_key returns (prefix_used, base, suffix). Rename to diffusers happens + in resolve_targets, not parse_key.""" cases = [ - # BFL prefix -> dotted base flattened to underscores ('diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight', C.LORA_SUFFIXES, - ('double_blocks_0_img_attn_proj', 'lora_down.weight')), - # PEFT prefix -> stays in diffusers form + ('diffusion_model.', 'double_blocks.0.img_attn.proj', 'lora_down.weight')), ('transformer.transformer_blocks.0.attn.to_q.lora_B.weight', C.LORA_SUFFIXES, - ('transformer_blocks_0_attn_to_q', 'lora_up.weight')), - # kohya prefix - already flat underscores + ('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_up.weight')), ('lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight', C.LORA_SUFFIXES, - ('double_blocks_0_img_attn_qkv', 'lora_down.weight')), - # Unrelated keys reject cleanly + ('lora_unet_', 'double_blocks_0_img_attn_qkv', 'lora_down.weight')), + # Bare BFL path (no prefix) + ('double_blocks.0.img_attn.proj.lora_A.weight', + C.LORA_SUFFIXES, + (None, 'double_blocks.0.img_attn.proj', 'lora_down.weight')), ('random.unrelated.key', C.LORA_SUFFIXES, None), ] for key, suffixes, expected in cases: @@ -575,20 +572,27 @@ def test_marker_disambiguation(): return True -def test_static_rename_table(): - """build_static_rename produces correct Flux to diffusers path remappings. - - Verifies the rename templates against the documented chroma_lora layout: - double_blocks_X_img_attn_proj -> transformer_blocks_X_attn_to_out_0, etc. +def test_resolve_targets_static_renames(): + """resolve_targets produces the documented Flux-to-diffusers remappings + for non-fused targets in both kohya and BFL forms. """ - rename = C.build_static_rename(N_DOUBLE, N_SINGLE) - # Spot-check key remappings - assert rename['double_blocks_0_img_attn_proj'] == 'transformer_blocks_0_attn_to_out_0' - assert rename['double_blocks_0_txt_attn_proj'] == 'transformer_blocks_0_attn_to_add_out' - assert rename['double_blocks_1_img_mlp_0'] == 'transformer_blocks_1_ff_net_0_proj' - assert rename['double_blocks_1_img_mlp_2'] == 'transformer_blocks_1_ff_net_2' - assert rename['double_blocks_0_txt_mlp_0'] == 'transformer_blocks_0_ff_context_net_0_proj' - assert rename['single_blocks_0_linear2'] == 'single_transformer_blocks_0_proj_out' + cases = [ + # kohya + (('lora_unet_', 'double_blocks_0_img_attn_proj'), 'transformer_blocks.0.attn.to_out.0'), + (('lora_unet_', 'double_blocks_0_txt_attn_proj'), 'transformer_blocks.0.attn.to_add_out'), + (('lora_unet_', 'double_blocks_1_img_mlp_0'), 'transformer_blocks.1.ff.net.0.proj'), + (('lora_unet_', 'double_blocks_1_img_mlp_2'), 'transformer_blocks.1.ff.net.2'), + (('lora_unet_', 'double_blocks_0_txt_mlp_0'), 'transformer_blocks.0.ff_context.net.0.proj'), + (('lora_unet_', 'single_blocks_0_linear2'), 'single_transformer_blocks.0.proj_out'), + # BFL dotted - same diffusers paths + (('diffusion_model.', 'double_blocks.0.img_attn.proj'), 'transformer_blocks.0.attn.to_out.0'), + (('diffusion_model.', 'single_blocks.0.linear2'), 'single_transformer_blocks.0.proj_out'), + ] + for (prefix, base), expected_path in cases: + targets = C.resolve_targets(prefix, base) + assert len(targets) == 1, f'({prefix}, {base}) -> {targets}' + path, chunk = targets[0] + assert path == expected_path and chunk is None, f'({prefix}, {base}) -> {targets}' return True @@ -747,13 +751,8 @@ def test_lokr_bfl_img_attn_proj(): return True -def test_lokr_bfl_img_attn_qkv_slice_chunked(): - """BFL LoKR on fused img_attn.qkv emits 3 SliceChunk modules with row ranges. - - Chroma's expand_chroma_fused_lokr uses NetworkModuleLokrSliceChunk for all - fused targets (even equal-chunks), since the implementation is general - and start/end ranges express both forms. - """ +def test_lokr_bfl_img_attn_qkv_chunked(): + """BFL LoKR on fused img_attn.qkv emits 3 LokrChunk modules (equal chunks).""" net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_qkv_equal_chunks()) assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}' expected = { @@ -763,9 +762,8 @@ def test_lokr_bfl_img_attn_qkv_slice_chunked(): } assert set(net.modules) == expected for nk, mod in net.modules.items(): - assert isinstance(mod, network_lokr.NetworkModuleLokrSliceChunk), f'{nk}: type={type(mod).__name__}' - # Each chunk row range is HIDDEN rows wide - assert mod.end_row - mod.start_row == HIDDEN, f'{nk}: range={mod.start_row}:{mod.end_row}' + assert isinstance(mod, network_lokr.NetworkModuleLokrChunk), f'{nk}: type={type(mod).__name__}' + assert mod.num_chunks == 3, f'{nk}: num_chunks={mod.num_chunks}' return True @@ -799,10 +797,18 @@ def test_loha_bfl_img_attn_proj(): return True -def test_loha_bfl_img_attn_qkv_skipped(): - """LoHA on fused img_attn.qkv is dropped (chroma has no LoHA chunk variant).""" +def test_loha_bfl_img_attn_qkv_chunked(): + """LoHA on fused img_attn.qkv emits 3 HadaChunk modules (equal chunks).""" net = _load_via(C.try_load_loha, sd_loha_bfl_img_attn_qkv_skipped()) - assert net is None or len(net.modules) == 0 + assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}' + expected = { + 'lora_transformer_transformer_blocks_0_attn_to_q', + 'lora_transformer_transformer_blocks_0_attn_to_k', + 'lora_transformer_transformer_blocks_0_attn_to_v', + } + assert set(net.modules) == expected + for mod in net.modules.values(): + assert isinstance(mod, network_hada.NetworkModuleHadaChunk) return True @@ -847,13 +853,13 @@ def test_lokr_calc_updown_shape(): return True -def test_lokr_slicechunk_equal_calc_updown_shape(): - """LokrSliceChunk with equal-width range produces (HIDDEN, HIDDEN) output.""" +def test_lokr_chunk_equal_calc_updown_shape(): + """LokrChunk equal-chunks dispatch produces (HIDDEN, HIDDEN) output for the QKV split.""" net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_qkv_equal_chunks()) mod = make_network_for_module(next(iter(net.modules.values()))) target = torch.randn(HIDDEN, HIDDEN) updown, _ = mod.calc_updown(target) - assert_shape(updown, target.shape, label='LokrSliceChunk equal range') + assert_shape(updown, target.shape, label='LokrChunk equal range') return True @@ -899,7 +905,7 @@ def run_tests(): t0 = time.time() log.warning('=== Parsing primitives ===') - for fn in [test_parse_key_all_prefixes, test_marker_disambiguation, test_static_rename_table]: + for fn in [test_parse_key_all_prefixes, test_marker_disambiguation, test_resolve_targets_static_renames]: run_test(CAT_PARSE, fn) log.warning('=== Loaders ===') @@ -917,10 +923,10 @@ def run_tests(): test_lora_distilled_guidance, test_lora_dora_threading, test_lokr_bfl_img_attn_proj, - test_lokr_bfl_img_attn_qkv_slice_chunked, + test_lokr_bfl_img_attn_qkv_chunked, test_lokr_bfl_single_linear1_unequal_chunks, test_loha_bfl_img_attn_proj, - test_loha_bfl_img_attn_qkv_skipped, + test_loha_bfl_img_attn_qkv_chunked, test_oft_bfl_img_attn_proj, test_oft_bfl_img_attn_qkv_skipped, ]: @@ -930,7 +936,7 @@ def run_tests(): for fn in [ test_lora_calc_updown_shape, test_lokr_calc_updown_shape, - test_lokr_slicechunk_equal_calc_updown_shape, + test_lokr_chunk_equal_calc_updown_shape, test_lokr_slicechunk_unequal_calc_updown_shape, test_loha_calc_updown_shape, test_oft_calc_updown_shape, From 6937ea803fb51c357dee3bf81e3eb1bdea88ee61 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 23:34:41 +0100 Subject: [PATCH 12/13] 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: From 4c3f4db1295df0cdadb659975a3ec144d023e425 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 23:38:37 +0100 Subject: [PATCH 13/13] refactor(lora): collapse native dispatcher to a registry Five if-blocks in lora_load.load_safetensors reduce to one lookup in _NATIVE_DISPATCH, a string -> module-path map keyed by shared.sd_model_type. Each entry's module exposes try_load(name, network_on_disk, lora_scale). flux2 / zimage / chroma / ernie use the umbrella that binds native_loader's generics via try_load_chain. Anima keeps its own try_load (aliased to try_load_lora) since its multi-component routing doesn't fit the shared suffix-table model. New native archs land by adding one entry to the dict and shipping a try_load. --- modules/lora/lora_load.py | 78 +++++++++-------------------------- pipelines/anima/anima_lora.py | 5 +++ 2 files changed, 25 insertions(+), 58 deletions(-) diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index b3c4b0908..d93386cab 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -17,6 +17,18 @@ exclude_errors = [ "'ChronoEditTransformer3DModel'", ] +# shared.sd_model_type -> dotted module path of a pipeline native loader +# exposing ``try_load(name, network_on_disk, lora_scale)``. New archs add an +# entry here and ship a per-arch ``try_load`` (either binding native_loader's +# generic helpers via try_load_chain, or rolling their own). +_NATIVE_DISPATCH = { + 'zimage': 'pipelines.z_image.zimage_lora', + 'chroma': 'pipelines.chroma.chroma_lora', + 'ernieimage': 'pipelines.ernie.ernie_lora', + 'f2': 'pipelines.flux.flux2_lora', + 'anima': 'pipelines.anima.anima_lora', +} + def lora_dump(lora, dct): import tempfile @@ -49,64 +61,14 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}') if cached is not None: return cached - if shared.sd_model_type == 'zimage': - from pipelines.z_image import zimage_lora - lora_scale = shared.opts.extra_networks_default_multiplier - zimage_net = None - for try_fn in (zimage_lora.try_load_lora, zimage_lora.try_load_lokr, zimage_lora.try_load_loha, zimage_lora.try_load_oft): - sub = try_fn(name, network_on_disk, lora_scale) - if sub is None: - continue - if zimage_net is None: - zimage_net = sub - else: - zimage_net.modules.update(sub.modules) - if zimage_net is not None: - lora_cache[name] = zimage_net - return zimage_net - if shared.sd_model_type == 'anima': - from pipelines.anima import anima_lora - lora_scale = shared.opts.extra_networks_default_multiplier - anima_net = anima_lora.try_load_lora(name, network_on_disk, lora_scale) - if anima_net is not None: - lora_cache[name] = anima_net - return anima_net - if shared.sd_model_type == 'ernieimage': - from pipelines.ernie import ernie_lora - lora_scale = shared.opts.extra_networks_default_multiplier - ernie_net = None - for try_fn in (ernie_lora.try_load_lora, ernie_lora.try_load_lokr, ernie_lora.try_load_loha, ernie_lora.try_load_oft): - sub = try_fn(name, network_on_disk, lora_scale) - if sub is None: - continue - if ernie_net is None: - ernie_net = sub - else: - ernie_net.modules.update(sub.modules) - if ernie_net is not None: - lora_cache[name] = ernie_net - return ernie_net - if shared.sd_model_type == 'chroma': - from pipelines.chroma import chroma_lora - lora_scale = shared.opts.extra_networks_default_multiplier - chroma_net = None - for try_fn in (chroma_lora.try_load_lora, chroma_lora.try_load_lokr, chroma_lora.try_load_loha, chroma_lora.try_load_oft): - sub = try_fn(name, network_on_disk, lora_scale) - if sub is None: - continue - if chroma_net is None: - chroma_net = sub - else: - chroma_net.modules.update(sub.modules) - if chroma_net is not None: - lora_cache[name] = chroma_net - return chroma_net - if shared.sd_model_type == 'f2': - from pipelines.flux import flux2_lora - f2_net = flux2_lora.try_load(name, network_on_disk, shared.opts.extra_networks_default_multiplier) - if f2_net is not None: - lora_cache[name] = f2_net - return f2_net + native_module = _NATIVE_DISPATCH.get(shared.sd_model_type) + if native_module is not None: + import importlib + mod = importlib.import_module(native_module) + net = mod.try_load(name, network_on_disk, shared.opts.extra_networks_default_multiplier) + if net is not None: + lora_cache[name] = net + return net net = network.Network(name, network_on_disk) net.mtime = os.path.getmtime(network_on_disk.filename) state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network') diff --git a/pipelines/anima/anima_lora.py b/pipelines/anima/anima_lora.py index 8454425de..dfa742ac2 100644 --- a/pipelines/anima/anima_lora.py +++ b/pipelines/anima/anima_lora.py @@ -63,6 +63,11 @@ COSMOS_2_FLAT_RENAME = OrderedDict([ ]) +def try_load(name, network_on_disk, lora_scale): + """Single dispatcher entry point. Anima only supports the LoRA family.""" + return try_load_lora(name, network_on_disk, lora_scale) + + def try_load_lora(name, network_on_disk, lora_scale): """Try loading an Anima LoRA as native modules.