From e45c23c0319c0cc1f1abaed83387e0c5d66ead98 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 18 May 2026 23:31:17 +0100 Subject: [PATCH] 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,