Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor(chroma): migrate to generic native_loader

Browse Source 
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.
This commit is contained in:
CalamitousFelicitousness
2026-05-18 23:31:17 +01:00
parent c3b7379fd3
commit e45c23c031
2 changed files with 241 additions and 495 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipelines/chroma/chroma_lora.py
+190 -450
View File
@@ -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}.<suffix> 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}.<suffix> 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),
)
test/test-chroma-native-adapters.py
+51 -45
View File
@@ -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,