feat(lora): native loading for qwen-image 2.1

QwenImage21Transformer2DModel keeps attention and the SwiGLU split, while ComfyUI checkpoints fuse the SwiGLU input as img_mlp.gate_up with rows [gate_layer; proj], the layout ai-toolkit and musubi-tuner train against. pipelines/qwen/qwen21_lora.py splits that projection by row halves for LoRA, DoRA, LoKr and LoHA in dotted and kohya-flattened form and binds every other key verbatim; qwen21 joins allow_native and NATIVE_DISPATCH. The diffusers path loads none of the fused files and fails on transformer.-prefixed files that carry alpha tensors. The tests map every layout found on CivitAI and the Hub against the real class, check the split against the fused SwiGLU forward, and refuse Qwen-Image 1.x files.

CHANGELOG.md

@@ -55,6 +55,7 @@ Plus quite a few other improvements and inevitable bug-fixes, see full changelog
 - **Other**
   - edit-models: simplify usage of multi-image reference editing  
     automatically handle: flux.2, glm-image, nano-banana, joy-image-edit, qwen-image-edit
+  - lora: native loading for qwen-image 2.1, including the fused mlp layout ai-toolkit and musubi-tuner train against
   - processors: `DWPose`, `RTMW`, `RTMO`: image processor updates, thanks @kirtasshh  
     refactored to use `ONNX` instead of obsolete `mmpose` lib and with additional configurable settings  
 - **UI**

modules/lora/lora_load.py

@@ -29,6 +29,7 @@ NATIVE_DISPATCH = {
     'anima':      'pipelines.anima.anima_lora',
     'krea2':      'pipelines.krea2.krea2_lora',
     'minimaxh3':  'pipelines.minimax.minimax_lora',
+    'qwen21':     'pipelines.qwen.qwen21_lora',
 }
 
 

modules/lora/lora_overrides.py

@@ -38,6 +38,7 @@ allow_native = [
     'ernieimage',
     'krea2',
     'minimaxh3',
+    'qwen21',
 ]
 
 

pipelines/qwen/qwen21_lora.py

@@ -0,0 +1,146 @@
+"""Qwen-Image 2.1 native adapter loader.
+
+Runs when :func:`modules.lora.lora_overrides.get_method` returns ``'native'``
+(``lora_force_diffusers`` off and ``qwen21`` in ``allow_native``).
+
+Entry points, one per family: :func:`try_load_lora` (plus DoRA),
+:func:`try_load_lokr`, :func:`try_load_loha`, :func:`try_load_oft`,
+:func:`try_load_ia3`, :func:`try_load_glora`, :func:`try_load_norm`,
+:func:`try_load_full`.
+
+``QwenImage21Transformer2DModel`` keeps attention split (``to_q`` / ``to_k`` /
+``to_v`` / ``to_out.0``) and the SwiGLU split (``img_mlp.gate_layer`` /
+``img_mlp.proj`` / ``img_mlp.out``), so diffusers-named keys bind verbatim
+under every prefix. ComfyUI checkpoints fuse the SwiGLU input as
+``img_mlp.gate_up`` with rows ``[gate_layer; proj]``, and so do the LoRAs
+ai-toolkit and musubi-tuner train against them; :func:`resolve_targets` splits
+that projection into the two diffusers modules.
+"""
+
+from modules.lora import native_adapter
+from modules.lora.native_adapter import ChunkSpec
+
+
+# === Arch-specific prefix configuration ===
+
+KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT
+
+
+# Fused SwiGLU input in dotted and kohya-flattened form, and the split modules its row halves land on.
+GATE_UP_SUFFIXES = (".img_mlp.gate_up", "_img_mlp_gate_up")
+GATE_UP_TARGETS = (("gate_layer", ChunkSpec(idx=0, total=2)), ("proj", ChunkSpec(idx=1, total=2)))
+
+
+# === Re-exports for test/back-compat ===
+
+LORA_SUFFIXES = native_adapter.LORA_SUFFIXES
+LOKR_SUFFIXES = native_adapter.LOKR_SUFFIXES
+LOHA_SUFFIXES = native_adapter.LOHA_SUFFIXES
+OFT_SUFFIXES = native_adapter.OFT_SUFFIXES
+IA3_SUFFIXES = native_adapter.IA3_SUFFIXES
+GLORA_SUFFIXES = native_adapter.GLORA_SUFFIXES
+NORM_SUFFIXES = native_adapter.NORM_SUFFIXES
+FULL_SUFFIXES = native_adapter.FULL_SUFFIXES
+
+LORA_MARKERS = native_adapter.LORA_MARKERS
+LOKR_MARKERS = native_adapter.LOKR_MARKERS
+LOHA_MARKERS = native_adapter.LOHA_MARKERS
+OFT_MARKERS = native_adapter.OFT_MARKERS
+IA3_MARKERS = native_adapter.IA3_MARKERS
+GLORA_MARKERS = native_adapter.GLORA_MARKERS
+NORM_MARKERS = native_adapter.NORM_MARKERS
+FULL_MARKERS = native_adapter.FULL_MARKERS
+
+SUFFIX_NORMALIZE = native_adapter.SUFFIX_NORMALIZE
+BARE_DIFFUSERS_PREFIX_USED = native_adapter.BARE_DIFFUSERS_PREFIX_USED
+has_marker = native_adapter.has_marker
+
+
+def parse_key(key, suffixes):
+    """Qwen-Image 2.1-bound :func:`native_adapter.parse_key`."""
+    return native_adapter.parse_key(
+        key, suffixes,
+        prefixes=KNOWN_PREFIXES,
+    )
+
+
+def group_by_suffixes(state_dict, suffixes):
+    """Qwen-Image 2.1-bound :func:`native_adapter.group_by_suffixes`."""
+    return native_adapter.group_by_suffixes(
+        state_dict, suffixes,
+        prefixes=KNOWN_PREFIXES,
+    )
+
+
+# === Target resolution (arch-specific) ===
+
+
+def resolve_targets(prefix_used, base):
+    """Return ``[(diffusers_path, ChunkSpec | None), ...]`` for a parsed group key.
+
+    A fused ``img_mlp.gate_up`` splits into ``gate_layer`` and ``proj`` under any
+    prefix, dotted or flattened. Everything else is already a diffusers path:
+    ``diffusion_model.`` and ``lora_unet_`` bind here, the universal passthrough
+    prefixes bind upstream in :func:`native_adapter.resolve_group_targets`.
+    """
+    for fused in GATE_UP_SUFFIXES:
+        if base.endswith(fused):
+            sep = fused[0]
+            stem = base[:-len(fused)]
+            return [(f"{stem}{sep}img_mlp{sep}{leaf}", chunk) for leaf, chunk in GATE_UP_TARGETS]
+    if prefix_used in ("diffusion_model.", "lora_unet_"):
+        return [(base, None)]
+    return []
+
+
+# === Native loaders (thin wrappers over native_adapter generics) ===
+
+
+BIND_KWARGS = dict(
+    resolve_targets=resolve_targets,
+    prefixes=KNOWN_PREFIXES,
+    arch_name="qwen21",
+)
+
+
+def try_load_lora(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_lora(name, network_on_disk, lora_scale, **BIND_KWARGS)
+
+
+def try_load_lokr(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_lokr(name, network_on_disk, lora_scale, **BIND_KWARGS)
+
+
+def try_load_loha(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_loha(name, network_on_disk, lora_scale, **BIND_KWARGS)
+
+
+def try_load_oft(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_oft(name, network_on_disk, lora_scale, **BIND_KWARGS)
+
+
+def try_load_ia3(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_ia3(name, network_on_disk, lora_scale, **BIND_KWARGS)
+
+
+def try_load_glora(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_glora(name, network_on_disk, lora_scale, **BIND_KWARGS)
+
+
+def try_load_norm(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_norm(name, network_on_disk, lora_scale, **BIND_KWARGS)
+
+
+def try_load_full(name, network_on_disk, lora_scale):
+    return native_adapter.try_load_full(name, network_on_disk, lora_scale, **BIND_KWARGS)
+
+
+def try_load(name, network_on_disk, lora_scale):
+    """Run every Qwen-Image 2.1 family loader, merge any that match."""
+    return native_adapter.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,
+        ),
+    )

test/test-qwen21-native-adapters.py

@@ -0,0 +1,693 @@
+#!/usr/bin/env python
+"""
+Offline unit tests for the Qwen-Image 2.1 native adapter loader.
+
+``diffusers.QwenImage21Transformer2DModel`` keeps attention and the SwiGLU
+split (``img_mlp.gate_layer`` / ``img_mlp.proj`` / ``img_mlp.out``), while
+ComfyUI checkpoints fuse the SwiGLU input as ``img_mlp.gate_up`` with rows
+``[gate_layer; proj]``, the layout ai-toolkit and musubi-tuner train against.
+``pipelines.qwen.qwen21_lora`` splits that projection; everything else binds
+verbatim. These tests pin the split (including a forward-pass check against
+the fused SwiGLU), every save layout seen on CivitAI and the Hub, and the
+file-level alpha read from PEFT / DiffSynth metadata.
+
+Save formats exercised:
+
+- ai-toolkit (``diffusion_model.transformer_blocks.0.img_mlp.gate_up.lora_A.weight``)
+- musubi-tuner (``lora_unet_transformer_blocks_0_img_mlp_gate_up.lora_down.weight`` + alpha)
+- diffusers-PEFT (``transformer.transformer_blocks.0.img_mlp.proj.lora_A.weight``),
+  PEFT named adapter (``...lora_A.default.weight``), bare factors (``...lora_A``)
+- naive diffusers-to-ComfyUI conversion (``diffusion_model.transformer.``)
+- LoKr and DoRA on the fused projection, rank-padded fused re-exports
+- Qwen-Image 1.x files, which must be refused
+
+The reference module tree is the real class at tiny dims. When the local
+``Lora/Qwen 2.1`` folder exists, every file in it is also mapped header-only
+against a full-size model built on the meta device.
+
+No running server required.
+
+Usage:
+    python test/test-qwen21-native-adapters.py
+"""
+
+import os
+import sys
+import math
+import json
+import tempfile
+
+import torch
+import torch.nn as nn
+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([])
+
+import diffusers  # pylint: disable=wrong-import-position
+from modules.errors import log   # pylint: disable=wrong-import-position
+from modules.lora import network, network_lora, network_lokr  # pylint: disable=wrong-import-position
+from modules.lora import native_adapter  # pylint: disable=wrong-import-position
+from modules.lora.native_adapter import ChunkSpec  # pylint: disable=wrong-import-position
+from pipelines.qwen import qwen21_lora as Q  # 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()
+        record(cat, ok is not False, 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()
+
+
+# ============================================================
+# Reference Qwen-Image 2.1 transformer (real class, tiny dims)
+# ============================================================
+# Upstream: 32 layers, 32 heads x 128 = 4096, mlp_ratio 3, context 4096, 64 latent channels.
+
+INNER = 32
+MLP = INNER * 3
+LAYERS = 2
+RANK = 4
+
+REF = diffusers.QwenImage21Transformer2DModel(
+    patch_size=1, in_channels=16, out_channels=16, num_layers=LAYERS,
+    attention_head_dim=16, num_attention_heads=2, context_in_dim=24,
+    mlp_ratio=3, axes_dims_rope=(4, 6, 6),
+)
+
+LINEAR_SHAPES = {name: tuple(m.weight.shape) for name, m in REF.named_modules() if isinstance(m, nn.Linear)}
+LINEAR_NETKEYS = {'lora_transformer_' + name.replace('.', '_') for name in LINEAR_SHAPES}
+BLOCK_LEAVES = ('attn.to_q', 'attn.to_k', 'attn.to_v', 'attn.to_out.0', 'img_mlp.gate_layer', 'img_mlp.proj', 'img_mlp.out')
+EXTRAS = ('modulation.1', 'time_text_embed.timestep_embedder.linear_1', 'time_text_embed.timestep_embedder.linear_2',
+          'txt_in.in_layer', 'txt_in.out_layer', 'img_in', 'norm_out.linear', 'proj_out')
+
+
+def shape(path):
+    assert path in LINEAR_SHAPES, f'reference has no Linear {path!r}'
+    return LINEAR_SHAPES[path]
+
+
+class MockPipeline:
+    def __init__(self, transformer):
+        self.transformer = transformer
+        self.text_encoder = None
+
+
+class MockSdModel:
+    def __init__(self, pipe):
+        self.pipe = pipe
+        self.network_layer_mapping = {}
+        self.embedding_db = None
+        self.__class__.__name__ = 'QwenImage21Pipeline'
+
+
+def install_mock_pipe(transformer=REF):
+    """Point shared.sd_model at a mock exposing ``transformer``; re-installed per load so stamps do not leak."""
+    from modules.modeldata import model_data
+    model_data.sd_model = MockSdModel(MockPipeline(transformer))
+    return model_data.sd_model
+
+
+# ============================================================
+# Synthesizers and file helpers
+# ============================================================
+
+
+def lora(key_base, out, inp, rank=RANK, names=('lora_A.weight', 'lora_B.weight'), alpha=None):
+    sd = {f'{key_base}.{names[0]}': torch.randn(rank, inp), f'{key_base}.{names[1]}': torch.randn(out, rank)}
+    if alpha is not None:
+        sd[f'{key_base}.alpha'] = torch.tensor(float(alpha))
+    return sd
+
+
+def ai_toolkit_block(i, prefix='diffusion_model.', names=('lora_A.weight', 'lora_B.weight'), alpha=None):
+    """One block as ai-toolkit saves it: split attention, fused gate_up, out."""
+    sd = {}
+    for leaf in ('attn.to_q', 'attn.to_k', 'attn.to_v', 'attn.to_out.0'):
+        sd.update(lora(f'{prefix}transformer_blocks.{i}.{leaf}', INNER, INNER, names=names, alpha=alpha))
+    sd.update(lora(f'{prefix}transformer_blocks.{i}.img_mlp.gate_up', 2 * MLP, INNER, names=names, alpha=alpha))
+    sd.update(lora(f'{prefix}transformer_blocks.{i}.img_mlp.out', INNER, MLP, names=names, alpha=alpha))
+    return sd
+
+
+def musubi_block(i, alpha=RANK):
+    """One block as musubi-tuner saves it: kohya-flattened names, lora_down/up, alpha."""
+    sd = {}
+    for leaf, (out, inp) in (('attn_to_q', (INNER, INNER)), ('attn_to_k', (INNER, INNER)), ('attn_to_v', (INNER, INNER)),
+                             ('attn_to_out_0', (INNER, INNER)), ('img_mlp_gate_up', (2 * MLP, INNER)), ('img_mlp_out', (INNER, MLP))):
+        sd.update(lora(f'lora_unet_transformer_blocks_{i}_{leaf}', out, inp, names=('lora_down.weight', 'lora_up.weight'), alpha=alpha))
+    return sd
+
+
+def split_block(i, prefix='transformer.', names=('lora_A.weight', 'lora_B.weight')):
+    """One block with diffusers names: split SwiGLU."""
+    sd = {}
+    for leaf in BLOCK_LEAVES:
+        out, inp = shape(f'transformer_blocks.{i}.{leaf}')
+        sd.update(lora(f'{prefix}transformer_blocks.{i}.{leaf}', out, inp, names=names))
+    return sd
+
+
+class TempLora:
+    """Context manager: writes a state dict (and optional metadata) to a temp safetensors file."""
+
+    def __init__(self, state_dict, name='test', metadata=None, cached_metadata=None):
+        self.state_dict = state_dict
+        self.name = name
+        self.metadata = metadata
+        self.cached_metadata = cached_metadata
+        self.path = None
+
+    def __enter__(self):
+        sd = {k: v.contiguous() 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, metadata=self.metadata)
+        return MockNetworkOnDisk(self.path, self.name, self.cached_metadata)
+
+    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, metadata=None):
+        self.filename = filename
+        self.name = name
+        self.shorthash = ''
+        self.sd_version = 'unknown'
+        if metadata is not None:
+            self.metadata = metadata
+
+
+def load_via(try_fn, state_dict, metadata=None, cached_metadata=None, name='test'):
+    install_mock_pipe()
+    with TempLora(state_dict, name=name, metadata=metadata, cached_metadata=cached_metadata) as nod:
+        net = try_fn(name, nod, lora_scale=1.0)
+    if net is not None:
+        for module in net.modules.values():
+            module.network.te_multiplier = 1.0
+            module.network.unet_multiplier = 1.0
+    return net
+
+
+def netkey(path):
+    return 'lora_transformer_' + path.replace('.', '_')
+
+
+def updown(module, path):
+    delta, _ = module.calc_updown(REF.get_submodule(path).weight.detach().clone())
+    return delta
+
+
+# ============================================================
+# Tests - parsing
+# ============================================================
+
+CAT_PARSE = category('parse')
+
+
+def test_parse_key_layouts():
+    bd = Q.BARE_DIFFUSERS_PREFIX_USED
+    cases = [
+        ('diffusion_model.transformer_blocks.0.img_mlp.gate_up.lora_A.weight', ('diffusion_model.', 'transformer_blocks.0.img_mlp.gate_up', 'lora_down.weight')),
+        ('transformer.transformer_blocks.0.attn.to_q.lora_B.weight', ('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_up.weight')),
+        ('diffusion_model.transformer.modulation.1.lora_A.weight', ('transformer.', 'modulation.1', 'lora_down.weight')),
+        ('transformer_blocks.0.attn.to_k.lora_A.default.weight', (bd, 'transformer_blocks.0.attn.to_k', 'lora_down.weight')),
+        ('transformer_blocks.0.attn.to_k.lora_B', (bd, 'transformer_blocks.0.attn.to_k', 'lora_up.weight')),
+        ('lora_unet_transformer_blocks_0_img_mlp_gate_up.lora_down.weight', ('lora_unet_', 'transformer_blocks_0_img_mlp_gate_up', 'lora_down.weight')),
+        ('lora_unet_transformer_blocks_0_attn_to_q.alpha', ('lora_unet_', 'transformer_blocks_0_attn_to_q', 'alpha')),
+        ('img_in.lora_down', (bd, 'img_in', 'lora_down.weight')),
+        ('diffusion_model.transformer_blocks.3.img_mlp.gate_up.dora_scale', ('diffusion_model.', 'transformer_blocks.3.img_mlp.gate_up', 'dora_scale')),
+        ('random.unrelated.key', None),
+    ]
+    for key, expected in cases:
+        got = Q.parse_key(key, Q.LORA_SUFFIXES)
+        assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
+    return True
+
+
+def test_markers():
+    assert Q.has_marker({'transformer_blocks.0.attn.to_q.lora_A': None}, Q.LORA_MARKERS), 'bare lora_A not detected'
+    assert Q.has_marker({'transformer_blocks.0.attn.to_q.lora_A.default.weight': None}, Q.LORA_MARKERS)
+    lokr = {'diffusion_model.transformer_blocks.0.img_mlp.gate_up.lokr_w1': None}
+    assert Q.has_marker(lokr, Q.LOKR_MARKERS) and not Q.has_marker(lokr, Q.LORA_MARKERS)
+    return True
+
+
+# ============================================================
+# Tests - target resolution
+# ============================================================
+
+CAT_RESOLVE = category('resolve')
+
+GATE_UP_SPLIT = [('transformer_blocks.0.img_mlp.gate_layer', ChunkSpec(idx=0, total=2)), ('transformer_blocks.0.img_mlp.proj', ChunkSpec(idx=1, total=2))]
+
+
+def test_resolve_gate_up_every_prefix():
+    for prefix in ('diffusion_model.', 'transformer.', Q.BARE_DIFFUSERS_PREFIX_USED, 'lora_transformer_'):
+        got = native_adapter.resolve_group_targets(Q.resolve_targets, prefix, 'transformer_blocks.0.img_mlp.gate_up')
+        assert got == GATE_UP_SPLIT, f'{prefix}: {got}'
+    return True
+
+
+def test_resolve_gate_up_flattened():
+    expected = [('transformer_blocks_0_img_mlp_gate_layer', ChunkSpec(idx=0, total=2)), ('transformer_blocks_0_img_mlp_proj', ChunkSpec(idx=1, total=2))]
+    for prefix in ('lora_unet_', 'lycoris_', 'lora_transformer_'):
+        got = native_adapter.resolve_group_targets(Q.resolve_targets, prefix, 'transformer_blocks_0_img_mlp_gate_up')
+        assert got == expected, f'{prefix}: {got}'
+        assert {netkey(path) for path, _ in got} <= LINEAR_NETKEYS
+    return True
+
+
+def test_resolve_verbatim_is_real():
+    for prefix in ('diffusion_model.', 'transformer.', Q.BARE_DIFFUSERS_PREFIX_USED):
+        for path in [f'transformer_blocks.1.{leaf}' for leaf in BLOCK_LEAVES] + list(EXTRAS):
+            got = native_adapter.resolve_group_targets(Q.resolve_targets, prefix, path)
+            assert got == [(path, None)], f'{prefix}{path}: {got}'
+            assert netkey(path) in LINEAR_NETKEYS, f'{path} is not a reference Linear'
+    got = native_adapter.resolve_group_targets(Q.resolve_targets, 'lora_unet_', 'transformer_blocks_1_attn_to_out_0')
+    assert got == [('transformer_blocks_1_attn_to_out_0', None)] and netkey('transformer_blocks.1.attn.to_out.0') in LINEAR_NETKEYS
+    return True
+
+
+# ============================================================
+# Tests - loading real-world layouts
+# ============================================================
+
+CAT_LOAD = category('load')
+FULL = 7 * LAYERS
+
+
+def test_load_ai_toolkit_fused():
+    sd = {}
+    for i in range(LAYERS):
+        sd.update(ai_toolkit_block(i))
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == FULL, f'modules={len(net.modules) if net else None}'
+    assert set(net.modules) <= LINEAR_NETKEYS
+    fused_up = sd['diffusion_model.transformer_blocks.0.img_mlp.gate_up.lora_B.weight']
+    down = sd['diffusion_model.transformer_blocks.0.img_mlp.gate_up.lora_A.weight']
+    gate = net.modules[netkey('transformer_blocks.0.img_mlp.gate_layer')]
+    proj = net.modules[netkey('transformer_blocks.0.img_mlp.proj')]
+    assert torch.equal(gate.up_model.weight, fused_up[:MLP]) and torch.equal(proj.up_model.weight, fused_up[MLP:]), 'gate_up rows landed on the wrong module'
+    assert torch.equal(gate.down_model.weight, down) and torch.equal(proj.down_model.weight, down)
+    return True
+
+
+def test_load_musubi_flattened_alpha():
+    sd = {}
+    for i in range(LAYERS):
+        sd.update(musubi_block(i, alpha=RANK / 2))
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == FULL, f'modules={len(net.modules) if net else None}'
+    for key in (netkey('transformer_blocks.1.img_mlp.gate_layer'), netkey('transformer_blocks.1.img_mlp.proj')):
+        assert net.modules[key].alpha == RANK / 2 and net.modules[key].calc_scale() == 0.5
+    return True
+
+
+def test_load_peft_named_adapter():
+    sd = {}
+    for i in range(LAYERS):
+        sd.update(split_block(i, prefix='', names=('lora_A.default.weight', 'lora_B.default.weight')))
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == FULL
+    return True
+
+
+def test_load_bare_factor_names():
+    sd = {}
+    for i in range(LAYERS):
+        sd.update(split_block(i, prefix='', names=('lora_A', 'lora_B')))
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == FULL, f'modules={len(net.modules) if net else None}'
+    return True
+
+
+def test_load_turbo_extras():
+    sd = {}
+    for i in range(LAYERS):
+        sd.update(split_block(i))
+    for path in ('modulation.1', 'time_text_embed.timestep_embedder.linear_1', 'time_text_embed.timestep_embedder.linear_2'):
+        sd.update(lora(f'transformer.{path}', *shape(path)))
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == FULL + 3
+    assert netkey('modulation.1') in net.modules
+    return True
+
+
+def test_load_comfy_wrapped_diffusers():
+    sd = {}
+    for i in range(LAYERS):
+        sd.update(split_block(i, prefix='diffusion_model.transformer.'))
+    sd.update(lora('diffusion_model.transformer.modulation.1', *shape('modulation.1')))
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == FULL + 1, f'modules={len(net.modules) if net else None}'
+    return True
+
+
+def test_load_pdd_backbone():
+    """alibaba-pai exports: bare names, bare lora_down/lora_up, non-block targets."""
+    sd = {}
+    for i in range(LAYERS):
+        sd.update(split_block(i, prefix='', names=('lora_down', 'lora_up')))
+    for path in ('img_in', 'txt_in.in_layer', 'txt_in.out_layer', 'norm_out.linear', 'modulation.1'):
+        sd.update(lora(path, *shape(path), names=('lora_down', 'lora_up')))
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == FULL + 5
+    return True
+
+
+def test_load_partial_blocks():
+    net = load_via(Q.try_load, ai_toolkit_block(1))
+    assert net is not None and len(net.modules) == 7
+    assert all('_blocks_1_' in key for key in net.modules)
+    return True
+
+
+def test_load_lokr_fused():
+    sd = {
+        'diffusion_model.transformer_blocks.0.img_mlp.gate_up.lokr_w1': torch.randn(8, 4),
+        'diffusion_model.transformer_blocks.0.img_mlp.gate_up.lokr_w2': torch.randn(2 * MLP // 8, INNER // 4),
+        'diffusion_model.transformer_blocks.0.img_mlp.gate_up.alpha': torch.tensor(4.0),
+    }
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == 2
+    gate = net.modules[netkey('transformer_blocks.0.img_mlp.gate_layer')]
+    proj = net.modules[netkey('transformer_blocks.0.img_mlp.proj')]
+    assert isinstance(gate, network_lokr.NetworkModuleLokrChunk) and (gate.chunk_index, gate.num_chunks) == (0, 2)
+    assert isinstance(proj, network_lokr.NetworkModuleLokrChunk) and (proj.chunk_index, proj.num_chunks) == (1, 2)
+    return True
+
+
+def test_load_dora_fused():
+    sd = lora('diffusion_model.transformer_blocks.0.img_mlp.gate_up', 2 * MLP, INNER)
+    sd['diffusion_model.transformer_blocks.0.img_mlp.gate_up.dora_scale'] = torch.rand(2 * MLP, 1) + 0.5
+    net = load_via(Q.try_load, sd)
+    assert net is not None and len(net.modules) == 2
+    ds = sd['diffusion_model.transformer_blocks.0.img_mlp.gate_up.dora_scale']
+    assert torch.equal(net.modules[netkey('transformer_blocks.0.img_mlp.gate_layer')].dora_scale, ds[:MLP])
+    assert torch.equal(net.modules[netkey('transformer_blocks.0.img_mlp.proj')].dora_scale, ds[MLP:])
+    return True
+
+
+def test_refuse_qwen_image_1():
+    """Qwen-Image 1.x names bind nothing; right names at the wrong width are a shape mismatch that refuses the file."""
+    names_only = lora('transformer.transformer_blocks.0.img_mod.1', 6 * INNER, INNER)
+    names_only.update(lora('transformer.transformer_blocks.0.txt_mlp.net.2', INNER, 4 * INNER))
+    assert load_via(Q.try_load, names_only) is None
+    wrong_width = lora('diffusion_model.transformer_blocks.0.attn.to_q', 3 * INNER, 3 * INNER)
+    wrong_width.update(lora('diffusion_model.transformer_blocks.0.attn.to_k', INNER, INNER))
+    assert load_via(Q.try_load, wrong_width) is None
+    return True
+
+
+# ============================================================
+# Tests - delta math
+# ============================================================
+
+CAT_MATH = category('math')
+
+
+def swiglu_fused(x, w_gate_up, w_out):
+    """ai-toolkit / ComfyUI forward: ``gate, up = gate_up(x).chunk(2)``."""
+    gate, up = (x @ w_gate_up.T).chunk(2, dim=-1)
+    return (torch.nn.functional.silu(gate) * up) @ w_out.T
+
+
+def swiglu_split(x, w_gate, w_proj, w_out):
+    """diffusers forward: ``out(silu(gate_layer(x)) * proj(x))``."""
+    return (torch.nn.functional.silu(x @ w_gate.T) * (x @ w_proj.T)) @ w_out.T
+
+
+def test_gate_up_split_matches_fused_forward():
+    """The split LoRA reproduces the fused SwiGLU output; the swapped split does not."""
+    mlp = REF.transformer_blocks[0].img_mlp
+    w_gate, w_proj, w_out = mlp.gate_layer.weight.detach(), mlp.proj.weight.detach(), mlp.out.weight.detach()
+    sd = lora('diffusion_model.transformer_blocks.0.img_mlp.gate_up', 2 * MLP, INNER)
+    delta = sd['diffusion_model.transformer_blocks.0.img_mlp.gate_up.lora_B.weight'] @ sd['diffusion_model.transformer_blocks.0.img_mlp.gate_up.lora_A.weight']
+    x = torch.randn(5, INNER)
+    reference = swiglu_fused(x, torch.cat([w_gate, w_proj]) + delta, w_out)
+    net = load_via(Q.try_load, sd)
+    d_gate = updown(net.modules[netkey('transformer_blocks.0.img_mlp.gate_layer')], 'transformer_blocks.0.img_mlp.gate_layer')
+    d_proj = updown(net.modules[netkey('transformer_blocks.0.img_mlp.proj')], 'transformer_blocks.0.img_mlp.proj')
+    native = swiglu_split(x, w_gate + d_gate, w_proj + d_proj, w_out)
+    assert torch.allclose(native, reference, atol=1e-5), f'max diff {(native - reference).abs().max():.3e}'
+    swapped = swiglu_split(x, w_gate + d_proj, w_proj + d_gate, w_out)
+    assert not torch.allclose(swapped, reference, atol=1e-3), 'swapped halves also match; the check does not discriminate'
+    return True
+
+
+def test_lokr_chunks_tile_the_kron():
+    w1, w2 = torch.randn(8, 4), torch.randn(2 * MLP // 8, INNER // 4)
+    sd = {
+        'diffusion_model.transformer_blocks.1.img_mlp.gate_up.lokr_w1': w1,
+        'diffusion_model.transformer_blocks.1.img_mlp.gate_up.lokr_w2': w2,
+    }
+    net = load_via(Q.try_load, sd)
+    full = torch.kron(w1, w2)
+    d_gate = updown(net.modules[netkey('transformer_blocks.1.img_mlp.gate_layer')], 'transformer_blocks.1.img_mlp.gate_layer')
+    d_proj = updown(net.modules[netkey('transformer_blocks.1.img_mlp.proj')], 'transformer_blocks.1.img_mlp.proj')
+    assert torch.allclose(torch.cat([d_gate, d_proj]), full, atol=1e-5)
+    return True
+
+
+def test_rank_padded_fused_is_exact():
+    """A split LoRA re-fused as ``down = [A_g; A_p]``, ``up = diag(B_g, B_p)`` returns each half exactly."""
+    a_g, b_g, a_p, b_p = torch.randn(RANK, INNER), torch.randn(MLP, RANK), torch.randn(RANK, INNER), torch.randn(MLP, RANK)
+    up = torch.zeros(2 * MLP, 2 * RANK)
+    up[:MLP, :RANK], up[MLP:, RANK:] = b_g, b_p
+    sd = {
+        'diffusion_model.transformer_blocks.0.img_mlp.gate_up.lora_down.weight': torch.cat([a_g, a_p]),
+        'diffusion_model.transformer_blocks.0.img_mlp.gate_up.lora_up.weight': up,
+        'diffusion_model.transformer_blocks.0.img_mlp.gate_up.alpha': torch.tensor(float(2 * RANK)),
+    }
+    net = load_via(Q.try_load, sd)
+    d_gate = updown(net.modules[netkey('transformer_blocks.0.img_mlp.gate_layer')], 'transformer_blocks.0.img_mlp.gate_layer')
+    d_proj = updown(net.modules[netkey('transformer_blocks.0.img_mlp.proj')], 'transformer_blocks.0.img_mlp.proj')
+    assert torch.allclose(d_gate, b_g @ a_g, atol=1e-5) and torch.allclose(d_proj, b_p @ a_p, atol=1e-5)
+    return True
+
+
+# ============================================================
+# Tests - file-level alpha from metadata
+# ============================================================
+
+CAT_ALPHA = category('alpha')
+
+
+def peft_metadata(**fields):
+    return {'lora_adapter_metadata': json.dumps({f'transformer.{k}': v for k, v in fields.items()}), 'format': 'pt'}
+
+
+def test_peft_metadata_alpha():
+    """Pruna ships r=64 lora_alpha=128 without alpha tensors: PEFT applies scale 2."""
+    net = load_via(Q.try_load, split_block(0), metadata=peft_metadata(r=RANK, lora_alpha=2 * RANK, alpha_pattern={}, use_rslora=False))
+    assert all(module.calc_scale() == 2.0 for module in net.modules.values())
+    return True
+
+
+def test_peft_alpha_pattern():
+    net = load_via(Q.try_load, split_block(0), metadata=peft_metadata(r=RANK, lora_alpha=RANK, alpha_pattern={'transformer_blocks.0.attn.to_q': RANK / 2}))
+    assert net.modules[netkey('transformer_blocks.0.attn.to_q')].calc_scale() == 0.5
+    assert net.modules[netkey('transformer_blocks.0.attn.to_k')].calc_scale() == 1.0
+    return True
+
+
+def test_peft_rslora():
+    net = load_via(Q.try_load, split_block(0), metadata=peft_metadata(r=RANK, lora_alpha=RANK, use_rslora=True))
+    assert all(math.isclose(module.calc_scale(), RANK / math.sqrt(RANK)) for module in net.modules.values())
+    return True
+
+
+def test_alpha_tensors_win_over_metadata():
+    sd = {}
+    for i in range(LAYERS):
+        sd.update(musubi_block(i, alpha=RANK))
+    net = load_via(Q.try_load, sd, metadata=peft_metadata(r=RANK, lora_alpha=4 * RANK))
+    assert all(module.calc_scale() == 1.0 for module in net.modules.values())
+    return True
+
+
+def test_flat_alpha_metadata():
+    """DiffSynth-Studio stores alpha and rank as flat metadata keys."""
+    net = load_via(Q.try_load, split_block(0, prefix='', names=('lora_A.default.weight', 'lora_B.default.weight')), metadata={'alpha': str(RANK / 2), 'rank': str(RANK)})
+    assert all(module.calc_scale() == 0.5 for module in net.modules.values())
+    return True
+
+
+def test_metadata_read_from_header_when_cache_empty():
+    """An empty cached metadata dict (failed read, --no-metadata) falls back to the file header."""
+    net = load_via(Q.try_load, split_block(0), metadata=peft_metadata(r=RANK, lora_alpha=2 * RANK), cached_metadata={})
+    assert all(module.calc_scale() == 2.0 for module in net.modules.values())
+    return True
+
+
+def test_text_encoder_component_is_separate():
+    components = native_adapter.peft_configs({'lora_adapter_metadata': json.dumps({'transformer.lora_alpha': 8, 'text_encoder.lora_alpha': 2})})
+    alpha = native_adapter.peft_alpha(components)
+    assert alpha('transformer.', 'transformer_blocks.0.attn.to_q', 4) == 8.0
+    assert alpha('lora_te_', 'model_layers_0_self_attn_q_proj', 4) == 2.0
+    return True
+
+
+# ============================================================
+# Tests - real files (header-only, full-size meta model)
+# ============================================================
+
+CAT_REAL = category('real-files')
+LORA_DIR = '/home/ohiom/database/models/Lora/Qwen 2.1'
+QWEN1_DIR = '/home/ohiom/database/models/Lora/Qwen'
+
+
+def full_size_shapes():
+    with torch.device('meta'):
+        model = diffusers.QwenImage21Transformer2DModel()
+    return {'lora_transformer_' + name.replace('.', '_'): tuple(m.weight.shape) for name, m in model.named_modules() if isinstance(m, nn.Linear)}
+
+
+def header_shapes(filename):
+    from safetensors import safe_open
+    with safe_open(filename, framework='pt', device='cpu') as f:
+        return {key: tuple(f.get_slice(key).get_shape()) for key in f.keys()}
+
+
+def map_file(shapes, targets):
+    """(mapped, unmapped, mismatch) for the LoRA, DoRA and LoKr groups of one file, resolved like the loaders do."""
+    mapped = unmapped = mismatch = 0
+    lora_groups = Q.group_by_suffixes(shapes, Q.LORA_SUFFIXES)
+    lokr_groups = Q.group_by_suffixes(shapes, Q.LOKR_SUFFIXES)
+    for (prefix, base), w in list(lora_groups.items()) + list(lokr_groups.items()):
+        is_lora = 'lora_down.weight' in w and 'lora_up.weight' in w
+        is_lokr = 'lokr_w1' in w and 'lokr_w2' in w
+        if not (is_lora or is_lokr):
+            continue
+        for path, chunk in native_adapter.resolve_group_targets(Q.resolve_targets, prefix, base):
+            key = 'lora_transformer_' + path.replace('.', '_')
+            if key not in targets:
+                unmapped += 1
+                continue
+            out, inp = targets[key]
+            total = chunk.total if chunk is not None else 1
+            if is_lora:
+                fits = w['lora_down.weight'][1] == inp and w['lora_up.weight'][0] == out * total
+            else:
+                fits = w['lokr_w1'][0] * w['lokr_w2'][0] == out * total and w['lokr_w1'][1] * w['lokr_w2'][1] == inp
+            mapped += fits
+            mismatch += not fits
+    return mapped, unmapped, mismatch
+
+
+def test_real_qwen21_files_map_cleanly():
+    if not os.path.isdir(LORA_DIR):
+        log.warning(f'  SKIP: {LORA_DIR} not present')
+        return True
+    targets = full_size_shapes()
+    files = sorted(f for f in os.listdir(LORA_DIR) if f.endswith('.safetensors'))
+    bad = []
+    for f in files:
+        mapped, unmapped, mismatch = map_file(header_shapes(os.path.join(LORA_DIR, f)), targets)
+        log.info(f'    {f}: mapped={mapped} unmapped={unmapped} mismatch={mismatch}')
+        if mapped == 0 or unmapped or mismatch:
+            bad.append(f)
+    assert not bad, f'{len(bad)}/{len(files)} files do not map cleanly: {bad}'
+    return True
+
+
+def test_real_qwen1_files_do_not_fit():
+    if not os.path.isdir(QWEN1_DIR):
+        log.warning(f'  SKIP: {QWEN1_DIR} not present')
+        return True
+    targets = full_size_shapes()
+    files = sorted(f for f in os.listdir(QWEN1_DIR) if f.endswith('.safetensors'))[:6]
+    for f in files:
+        mapped, unmapped, mismatch = map_file(header_shapes(os.path.join(QWEN1_DIR, f)), targets)
+        log.info(f'    {f}: mapped={mapped} unmapped={unmapped} mismatch={mismatch}')
+        assert unmapped or mismatch, f'{f} maps onto Qwen-Image 2.1 without a single mismatch'
+    return True
+
+
+# ============================================================
+# Main
+# ============================================================
+
+
+def main():
+    tests = {
+        CAT_PARSE: [test_parse_key_layouts, test_markers],
+        CAT_RESOLVE: [test_resolve_gate_up_every_prefix, test_resolve_gate_up_flattened, test_resolve_verbatim_is_real],
+        CAT_LOAD: [
+            test_load_ai_toolkit_fused, test_load_musubi_flattened_alpha, test_load_peft_named_adapter,
+            test_load_bare_factor_names, test_load_turbo_extras, test_load_comfy_wrapped_diffusers,
+            test_load_pdd_backbone, test_load_partial_blocks, test_load_lokr_fused, test_load_dora_fused,
+            test_refuse_qwen_image_1,
+        ],
+        CAT_MATH: [test_gate_up_split_matches_fused_forward, test_lokr_chunks_tile_the_kron, test_rank_padded_fused_is_exact],
+        CAT_ALPHA: [
+            test_peft_metadata_alpha, test_peft_alpha_pattern, test_peft_rslora, test_alpha_tensors_win_over_metadata,
+            test_flat_alpha_metadata, test_metadata_read_from_header_when_cache_empty, test_text_encoder_component_is_separate,
+        ],
+        CAT_REAL: [test_real_qwen21_files_map_cleanly, test_real_qwen1_files_do_not_fit],
+    }
+    torch.manual_seed(0)
+    for cat, fns in tests.items():
+        log.info(f'Category: {cat}')
+        for fn in fns:
+            run_test(cat, fn)
+    passed = sum(r['passed'] for r in results.values())
+    failed = sum(r['failed'] for r in results.values())
+    for cat, r in results.items():
+        log.info(f'{cat}: passed={r["passed"]} failed={r["failed"]}')
+    log.info(f'Total: passed={passed} failed={failed}')
+    return 1 if failed else 0
+
+
+if __name__ == '__main__':
+    sys.exit(main())