mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 08:44:33 +02:00
CalamitousFelicitousness
b434eb0c1b
The parser no longer takes reference-name prefixes to tell bare reference keys from bare diffusers keys. Any bare key carries the sentinel and the arch resolver renames what it knows and passes the rest through. Flux2 keeps its list for file-format detection only.
1669 lines
73 KiB
Python
1669 lines
73 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
Offline unit tests for Flux2/Klein native adapter loaders.
|
|
|
|
Covers the native families (LoRA, LoKR, LoHA, OFT, BOFT, IA3, GLoRA,
|
|
Norm, Full) plus DoRA threading via the universal
|
|
NetworkModule.finalize_updown hook, and ex_bias accumulation across
|
|
stacked Norm adapters.
|
|
|
|
The tests build a mock Flux2-shaped transformer, write synthetic
|
|
safetensors files for each adapter format, and exercise the full loader
|
|
path: state-dict -> group_by_suffixes -> resolve_targets ->
|
|
NetworkModule* -> calc_updown -> finalize_updown.
|
|
|
|
Loader correctness is verified against the documented LyCORIS / kohya
|
|
save formats. End-to-end inference verification against real adapter
|
|
files remains an open gap for the rarer families.
|
|
|
|
No running server required.
|
|
|
|
Usage:
|
|
python test/test-flux2-native-adapters.py
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
|
|
import torch
|
|
import safetensors.torch
|
|
|
|
script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, script_dir)
|
|
os.chdir(script_dir)
|
|
|
|
os.environ['SD_INSTALL_QUIET'] = '1'
|
|
|
|
# Bootstrap cmd_args before any module that pulls in shared.py.
|
|
# parse_args() registers the full main+compat option set; installer.add_args
|
|
# adds the installer-specific options on top.
|
|
import modules.cmd_args # pylint: disable=wrong-import-position
|
|
import installer # pylint: disable=wrong-import-position
|
|
_orig_argv = sys.argv
|
|
sys.argv = [sys.argv[0]]
|
|
try:
|
|
modules.cmd_args.parse_args()
|
|
finally:
|
|
sys.argv = _orig_argv
|
|
installer.add_args(modules.cmd_args.parser)
|
|
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
|
|
|
|
from modules.logger import log # pylint: disable=wrong-import-position
|
|
from modules import shared # pylint: disable=wrong-import-position
|
|
from modules.lora import ( # pylint: disable=wrong-import-position
|
|
network, network_lora, network_lokr, network_hada, network_oft, network_boft,
|
|
network_ia3, network_glora, network_norm, network_full,
|
|
)
|
|
from modules.lora import lora_apply # pylint: disable=wrong-import-position
|
|
from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position
|
|
from pipelines.flux import flux2_lora as F # pylint: disable=wrong-import-position
|
|
|
|
|
|
# ============================================================
|
|
# Test infrastructure
|
|
# ============================================================
|
|
|
|
results: dict[str, dict] = {}
|
|
|
|
|
|
def category(name: str):
|
|
if name not in results:
|
|
results[name] = {'passed': 0, 'failed': 0, 'tests': []}
|
|
return name
|
|
|
|
|
|
def record(cat: str, passed: bool, name: str, detail: str = ''):
|
|
status = 'PASS' if passed else 'FAIL'
|
|
results[cat]['passed' if passed else 'failed'] += 1
|
|
results[cat]['tests'].append((status, name))
|
|
msg = f' {status}: {name}'
|
|
if detail:
|
|
msg += f' ({detail})'
|
|
if passed:
|
|
log.info(msg)
|
|
else:
|
|
log.error(msg)
|
|
|
|
|
|
def run_test(cat: str, fn):
|
|
name = fn.__name__
|
|
try:
|
|
ok = fn()
|
|
if ok is False:
|
|
record(cat, False, name)
|
|
else:
|
|
record(cat, True, name)
|
|
except AssertionError as e:
|
|
record(cat, False, name, str(e))
|
|
except Exception as e: # pylint: disable=broad-except
|
|
record(cat, False, name, f'exception: {e}')
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
|
|
# ============================================================
|
|
# Mock Flux2-shaped transformer
|
|
# ============================================================
|
|
|
|
# Small but realistic shapes so tests are fast.
|
|
HIDDEN = 128
|
|
QKV_OUT = 128
|
|
MLP_OUT = 512
|
|
HEAD_DIM = 32
|
|
SINGLE_FUSED_OUT = 3 * QKV_OUT + MLP_OUT # to_qkv_mlp_proj
|
|
SINGLE_OUT_IN = QKV_OUT + MLP_OUT # to_out
|
|
N_DOUBLE = 2
|
|
N_SINGLE = 2
|
|
|
|
|
|
# pylint: disable=attribute-defined-outside-init
|
|
# `_Holder` is a generic container we attach diffusers-shaped children to at
|
|
# runtime to mirror Flux2's transformer module tree.
|
|
class _Holder(torch.nn.Module):
|
|
"""Empty container module — we attach Linear/RMSNorm children dynamically."""
|
|
|
|
|
|
def build_mock_transformer():
|
|
"""Build a torch.nn.Module mimicking Flux2's diffusers-side module tree.
|
|
|
|
Mirrors the paths in F2_SINGLE_MAP / F2_DOUBLE_MAP / F2_QKV_MAP /
|
|
F2_EXTRA_MAP plus the RMSNorm targets inside attention so that
|
|
try_load_norm has something to bind. Sizes are scaled-down but
|
|
proportional so chunking math is realistic.
|
|
"""
|
|
transformer = _Holder()
|
|
transformer.transformer_blocks = torch.nn.ModuleList()
|
|
for _ in range(N_DOUBLE):
|
|
block = _Holder()
|
|
block.attn = _Holder()
|
|
block.attn.to_q = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
|
|
block.attn.to_k = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
|
|
block.attn.to_v = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
|
|
block.attn.to_out = torch.nn.ModuleList([torch.nn.Linear(QKV_OUT, HIDDEN, bias=False)])
|
|
block.attn.add_q_proj = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
|
|
block.attn.add_k_proj = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
|
|
block.attn.add_v_proj = torch.nn.Linear(HIDDEN, QKV_OUT, bias=False)
|
|
block.attn.to_add_out = torch.nn.Linear(QKV_OUT, HIDDEN, bias=False)
|
|
block.attn.norm_q = torch.nn.RMSNorm(HEAD_DIM)
|
|
block.attn.norm_k = torch.nn.RMSNorm(HEAD_DIM)
|
|
block.attn.norm_added_q = torch.nn.RMSNorm(HEAD_DIM)
|
|
block.attn.norm_added_k = torch.nn.RMSNorm(HEAD_DIM)
|
|
block.ff = _Holder()
|
|
block.ff.linear_in = torch.nn.Linear(HIDDEN, MLP_OUT, bias=False)
|
|
block.ff.linear_out = torch.nn.Linear(MLP_OUT, HIDDEN, bias=False)
|
|
block.ff_context = _Holder()
|
|
block.ff_context.linear_in = torch.nn.Linear(HIDDEN, MLP_OUT, bias=False)
|
|
block.ff_context.linear_out = torch.nn.Linear(MLP_OUT, HIDDEN, bias=False)
|
|
transformer.transformer_blocks.append(block)
|
|
transformer.single_transformer_blocks = torch.nn.ModuleList()
|
|
for _ in range(N_SINGLE):
|
|
sblock = _Holder()
|
|
sblock.attn = _Holder()
|
|
sblock.attn.to_qkv_mlp_proj = torch.nn.Linear(HIDDEN, SINGLE_FUSED_OUT, bias=False)
|
|
sblock.attn.to_out = torch.nn.Linear(SINGLE_OUT_IN, HIDDEN, bias=False)
|
|
transformer.single_transformer_blocks.append(sblock)
|
|
# Non-block F2_EXTRA_MAP targets. Modulation/norm_out out-dims follow the
|
|
# real modules' expansion ratios (3x/6x per mod_param_sets, 2x for AdaLN).
|
|
transformer.x_embedder = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
|
|
transformer.context_embedder = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
|
|
transformer.time_guidance_embed = _Holder()
|
|
transformer.time_guidance_embed.timestep_embedder = _Holder()
|
|
transformer.time_guidance_embed.timestep_embedder.linear_1 = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
|
|
transformer.time_guidance_embed.timestep_embedder.linear_2 = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
|
|
transformer.time_guidance_embed.guidance_embedder = _Holder()
|
|
transformer.time_guidance_embed.guidance_embedder.linear_1 = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
|
|
transformer.time_guidance_embed.guidance_embedder.linear_2 = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
|
|
transformer.norm_out = _Holder()
|
|
transformer.norm_out.linear = torch.nn.Linear(HIDDEN, 2 * HIDDEN, bias=False)
|
|
transformer.proj_out = torch.nn.Linear(HIDDEN, HIDDEN, bias=False)
|
|
transformer.single_stream_modulation = _Holder()
|
|
transformer.single_stream_modulation.linear = torch.nn.Linear(HIDDEN, 3 * HIDDEN, bias=False)
|
|
transformer.double_stream_modulation_img = _Holder()
|
|
transformer.double_stream_modulation_img.linear = torch.nn.Linear(HIDDEN, 6 * HIDDEN, bias=False)
|
|
transformer.double_stream_modulation_txt = _Holder()
|
|
transformer.double_stream_modulation_txt.linear = torch.nn.Linear(HIDDEN, 6 * HIDDEN, bias=False)
|
|
return transformer
|
|
|
|
|
|
class _MockFlux2Pipeline:
|
|
"""Class name carries 'Flux2' so shared.sd_model_type returns 'f2'
|
|
via modeldata.get_model_type's name-based dispatch."""
|
|
|
|
def __init__(self, transformer):
|
|
self.transformer = transformer
|
|
self.text_encoder = None
|
|
|
|
|
|
class _MockFlux2SdModel:
|
|
"""Outer wrapper holding the pipe + the network_layer_mapping that
|
|
lora_convert.assign_network_names_to_compvis_modules writes onto."""
|
|
|
|
def __init__(self, pipe):
|
|
self.pipe = pipe
|
|
self.network_layer_mapping = {}
|
|
self.embedding_db = None
|
|
self.__class__.__name__ = 'Flux2Pipeline' # belt-and-suspenders
|
|
|
|
|
|
def install_mock_pipe():
|
|
"""Set shared.sd_model to a mock exposing a Flux2-shaped transformer.
|
|
|
|
Each test calls this fresh so module.network_layer_name stamps from prior
|
|
tests don't leak (notably the loader-local stamping in try_load_norm).
|
|
|
|
Bypasses the ModelData lock by writing directly to model_data.sd_model;
|
|
shared.sd_model = ... goes through set_sd_model which is a no-op when
|
|
model_data.locked is True (the default outside webui startup).
|
|
"""
|
|
transformer = build_mock_transformer()
|
|
pipe = _MockFlux2Pipeline(transformer)
|
|
sd_model = _MockFlux2SdModel(pipe)
|
|
from modules.modeldata import model_data
|
|
model_data.sd_model = sd_model
|
|
return sd_model
|
|
|
|
|
|
# ============================================================
|
|
# State-dict synthesizers (one per family)
|
|
# ============================================================
|
|
|
|
# Realistic ranks for the synthesized factors.
|
|
RANK_LORA = 8
|
|
RANK_LOKR = 4
|
|
LOKR_W1_DIM = 8
|
|
LOKR_W2_DIM = HIDDEN // LOKR_W1_DIM
|
|
|
|
|
|
def sd_lora_kohya_qkv():
|
|
"""Kohya-format LoRA targeting fused img_attn.qkv (3-way chunked at load)."""
|
|
return {
|
|
# Down weight is shared across Q/K/V; up weight is fused (3*QKV_OUT, RANK).
|
|
'lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight': torch.randn(3 * QKV_OUT, RANK_LORA),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)),
|
|
}
|
|
|
|
|
|
def sd_lora_bfl_proj():
|
|
"""BFL/AI-toolkit-format LoRA on a non-fused double_blocks proj target."""
|
|
return {
|
|
'diffusion_model.double_blocks.1.img_attn.proj.lora_A.weight': torch.randn(RANK_LORA, QKV_OUT),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
|
|
}
|
|
|
|
|
|
def sd_lora_peft_to_q():
|
|
"""Diffusers PEFT-format LoRA targeting a single split-QKV diffusers path."""
|
|
return {
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(QKV_OUT, RANK_LORA),
|
|
}
|
|
|
|
|
|
def sd_lora_bare_bfl_mlp():
|
|
"""Bare BFL keys (no prefix) on an MLP target — exercises the bare fallback."""
|
|
return {
|
|
'double_blocks.0.img_mlp.0.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'double_blocks.0.img_mlp.0.lora_B.weight': torch.randn(MLP_OUT, RANK_LORA),
|
|
}
|
|
|
|
|
|
def sd_lora_peft_saved_fal_style():
|
|
"""PEFT-saved LoRA: base_model.model. wrapper around BFL paths.
|
|
|
|
Mirrors the actual key layout of fal/flux-2-klein-4B-outpaint-lora as
|
|
inspected from its safetensors header on HF — first 12 keys all of the
|
|
form ``base_model.model.double_blocks.0.img_attn.proj.lora_A.weight``.
|
|
"""
|
|
return {
|
|
'base_model.model.double_blocks.1.img_attn.proj.lora_A.weight': torch.randn(RANK_LORA, QKV_OUT),
|
|
'base_model.model.double_blocks.1.img_attn.proj.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
|
|
'base_model.model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
|
|
}
|
|
|
|
|
|
def sd_lora_peft_saved_dreambooth_style():
|
|
"""PEFT-saved LoRA where the wrapped path is diffusers-form (HF DreamBooth).
|
|
|
|
HF DreamBooth scripts target diffusers modules by name (e.g.
|
|
``transformer.transformer_blocks.0.attn.to_q``); after peft.save_pretrained
|
|
wraps the in-memory keys, the saved form is
|
|
``base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.weight``.
|
|
"""
|
|
return {
|
|
'base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(QKV_OUT, RANK_LORA),
|
|
}
|
|
|
|
|
|
def sd_lora_peft_saved_diffusers_style():
|
|
"""PEFT-saved LoRA with bare-diffusers paths and a ``.default.`` adapter-name infix.
|
|
|
|
Mirrors the actual key layout of community Flux2 LoRAs exported via
|
|
``Flux2Transformer2DModel.save_lora_adapter()`` after attaching an unnamed
|
|
PEFT adapter (e.g. ``lenovo_flux_klein9b.safetensors``):
|
|
|
|
- No ``transformer.`` / ``diffusion_model.`` / ``base_model.model.`` wrapper
|
|
- Keys start with bare ``transformer_blocks.`` or ``single_transformer_blocks.``
|
|
- ``.lora_A.default.weight`` / ``.lora_B.default.weight`` (PEFT adapter slot)
|
|
"""
|
|
return {
|
|
# Single-block fused module — exercises the bare-diffusers prefix branch.
|
|
'single_transformer_blocks.0.attn.to_qkv_mlp_proj.lora_A.default.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'single_transformer_blocks.0.attn.to_qkv_mlp_proj.lora_B.default.weight': torch.randn(SINGLE_FUSED_OUT, RANK_LORA),
|
|
# Double-block split target — exercises both the bare-diffusers branch and the .default. strip.
|
|
'transformer_blocks.0.attn.to_q.lora_A.default.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'transformer_blocks.0.attn.to_q.lora_B.default.weight': torch.randn(QKV_OUT, RANK_LORA),
|
|
}
|
|
|
|
|
|
def sd_lora_with_dora_scale():
|
|
"""LoRA with a dora_scale companion vector — exercises DoRA threading."""
|
|
return {
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(QKV_OUT, RANK_LORA),
|
|
'transformer.transformer_blocks.0.attn.to_q.dora_scale': torch.randn(QKV_OUT),
|
|
}
|
|
|
|
|
|
def sd_lora_dora_fused_qkv():
|
|
"""Kohya LoRA on fused img_attn.qkv with per-output dora_scale (LyCORIS wd_on_out=True)."""
|
|
return {
|
|
'lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight': torch.randn(3 * QKV_OUT, RANK_LORA),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.dora_scale': torch.rand(3 * QKV_OUT, 1) + 0.5,
|
|
}
|
|
|
|
|
|
def sd_lora_diff_b_fused_qkv():
|
|
"""Kohya LoRA on fused img_attn.qkv with a per-output diff_b bias delta."""
|
|
sd = dict(sd_lora_dora_fused_qkv())
|
|
del sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale']
|
|
sd['lora_unet_double_blocks_0_img_attn_qkv.diff_b'] = torch.randn(3 * QKV_OUT)
|
|
return sd
|
|
|
|
|
|
def sd_lora_aitk_magnitude_dora(w_base):
|
|
"""ai-toolkit DoRA save: lora_A/B + 1-D magnitude row norms, no alpha key.
|
|
|
|
Targets img_attn.proj -> attn.to_out.0, a square Linear (HIDDEN == QKV_OUT
|
|
in the mock): the case where a raw 1-D vector cannot be disambiguated
|
|
between per-output and per-input at apply time.
|
|
"""
|
|
down = torch.randn(RANK_LORA, QKV_OUT)
|
|
up = torch.randn(HIDDEN, RANK_LORA)
|
|
magnitude = w_base.reshape(HIDDEN, -1).norm(dim=1) # 1-D (out,)
|
|
return {
|
|
'diffusion_model.double_blocks.1.img_attn.proj.lora_A.weight': down,
|
|
'diffusion_model.double_blocks.1.img_attn.proj.lora_B.weight': up,
|
|
'diffusion_model.double_blocks.1.img_attn.proj.magnitude': magnitude,
|
|
}, down, up, magnitude
|
|
|
|
|
|
def sd_lokr_dora_fused_qkv(per_input=False):
|
|
"""Kohya LoKR on fused img_attn.qkv with a dora_scale companion.
|
|
|
|
``per_input=True`` stores the wd_on_out=False orientation ``(1, in)``,
|
|
which has no exact per-chunk split and must be skipped by the loader.
|
|
"""
|
|
sd = dict(sd_lokr_kohya_qkv())
|
|
if per_input:
|
|
ds = torch.rand(1, HIDDEN) + 0.5
|
|
else:
|
|
ds = torch.rand(3 * QKV_OUT, 1) + 0.5
|
|
sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale'] = ds
|
|
return sd
|
|
|
|
|
|
def sd_lokr_bfl_proj():
|
|
"""BFL-format LoKR on a non-fused proj target."""
|
|
return {
|
|
# w1 = (LOKR_W1_DIM, LOKR_W1_DIM), w2 = (LOKR_W2_DIM, LOKR_W2_DIM); kron is (HIDDEN, HIDDEN).
|
|
# Target shape is (HIDDEN, QKV_OUT) here; we need shape compat for calc_updown.
|
|
# Use w1=(QKV_OUT/LOKR_W2_DIM, HIDDEN/LOKR_W2_DIM) outer factor; LyCORIS picks dims by factorization.
|
|
'diffusion_model.double_blocks.1.img_attn.proj.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.lokr_w2': torch.randn(QKV_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(LOKR_W1_DIM)),
|
|
}
|
|
|
|
|
|
def sd_lokr_kohya_qkv():
|
|
"""Kohya LoKR targeting fused img_attn.qkv — exercises NetworkModuleLokrChunk."""
|
|
# Target after kron is (3*QKV_OUT, HIDDEN); chunk along dim=0 gives Q/K/V slices.
|
|
fused_out = 3 * QKV_OUT
|
|
# Pick w1=(8, ...) so kron of w1 with w2 spans fused_out x HIDDEN.
|
|
w1_dim_out = 8
|
|
w2_dim_out = fused_out // w1_dim_out
|
|
w1_dim_in = 4
|
|
w2_dim_in = HIDDEN // w1_dim_in
|
|
return {
|
|
'lora_unet_double_blocks_0_img_attn_qkv.lokr_w1': torch.randn(w1_dim_out, w1_dim_in),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.lokr_w2': torch.randn(w2_dim_out, w2_dim_in),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(w1_dim_in)),
|
|
}
|
|
|
|
|
|
def sd_lokr_simpletuner_lycoris_style():
|
|
"""LyCORIS-standalone-format LoKR (SimpleTuner save).
|
|
|
|
Mirrors the actual key layout of markury/flux2k9b-simpletuner-lokr-loona
|
|
as inspected from its safetensors header on HF — keys of the form
|
|
``lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w1`` (288 such keys
|
|
in the real file). The path under ``lycoris_`` is a diffusers path with
|
|
dots rendered as underscores; resolve_targets returns it verbatim and
|
|
the caller's ``.replace('.', '_')`` is a no-op.
|
|
|
|
Two targets here exercise both an img-side projection (to_q) and a
|
|
txt-side projection (add_k_proj) which uses underscores in the module
|
|
name itself — the latter is the case where naïve underscore-to-dot
|
|
expansion would corrupt the path.
|
|
"""
|
|
return {
|
|
'lycoris_transformer_blocks_0_attn_to_q.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
|
|
'lycoris_transformer_blocks_0_attn_to_q.lokr_w2': torch.randn(QKV_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
|
|
'lycoris_transformer_blocks_0_attn_to_q.alpha': torch.tensor(float(LOKR_W1_DIM)),
|
|
'lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
|
|
'lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w2': torch.randn(QKV_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
|
|
'lycoris_transformer_blocks_0_attn_add_k_proj.alpha': torch.tensor(float(LOKR_W1_DIM)),
|
|
}
|
|
|
|
|
|
def sd_lokr_bfl_extra_modules():
|
|
"""BFL-format LoKR spanning every non-block target in F2_EXTRA_MAP.
|
|
|
|
Mirrors the ai-toolkit full-matrix layout (klein_snofs_v1_1): both factors
|
|
stored whole (lokr_w1/lokr_w2, no a/b decomposition) with the ~1e10
|
|
placeholder alpha. With no rank-decomposed factor the loader must leave
|
|
dim=None so the placeholder never scales the delta.
|
|
"""
|
|
dims = {
|
|
'img_in': (HIDDEN, HIDDEN),
|
|
'txt_in': (HIDDEN, HIDDEN),
|
|
'time_in.in_layer': (HIDDEN, HIDDEN),
|
|
'time_in.out_layer': (HIDDEN, HIDDEN),
|
|
'guidance_in.in_layer': (HIDDEN, HIDDEN),
|
|
'guidance_in.out_layer': (HIDDEN, HIDDEN),
|
|
'final_layer.linear': (HIDDEN, HIDDEN),
|
|
'final_layer.adaLN_modulation.1': (2 * HIDDEN, HIDDEN),
|
|
'single_stream_modulation.lin': (3 * HIDDEN, HIDDEN),
|
|
'double_stream_modulation_img.lin': (6 * HIDDEN, HIDDEN),
|
|
'double_stream_modulation_txt.lin': (6 * HIDDEN, HIDDEN),
|
|
}
|
|
sd = {}
|
|
for base, (out_dim, in_dim) in dims.items():
|
|
sd[f'diffusion_model.{base}.lokr_w1'] = torch.randn(LOKR_W1_DIM, LOKR_W1_DIM)
|
|
sd[f'diffusion_model.{base}.lokr_w2'] = torch.randn(out_dim // LOKR_W1_DIM, in_dim // LOKR_W1_DIM)
|
|
sd[f'diffusion_model.{base}.alpha'] = torch.tensor(9999220736.0)
|
|
return sd
|
|
|
|
|
|
def sd_loha_bfl_proj():
|
|
"""BFL-format LoHA on a non-fused proj target."""
|
|
return {
|
|
'diffusion_model.double_blocks.1.img_attn.proj.hada_w1_a': torch.randn(HIDDEN, RANK_LORA),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.hada_w1_b': torch.randn(RANK_LORA, QKV_OUT),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.hada_w2_a': torch.randn(HIDDEN, RANK_LORA),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.hada_w2_b': torch.randn(RANK_LORA, QKV_OUT),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
|
|
}
|
|
|
|
|
|
def sd_loha_kohya_qkv():
|
|
"""Kohya LoHA targeting fused img_attn.qkv — exercises NetworkModuleHadaChunk."""
|
|
fused_out = 3 * QKV_OUT
|
|
return {
|
|
'lora_unet_double_blocks_0_img_attn_qkv.hada_w1_a': torch.randn(fused_out, RANK_LORA),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.hada_w1_b': torch.randn(RANK_LORA, HIDDEN),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.hada_w2_a': torch.randn(fused_out, RANK_LORA),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.hada_w2_b': torch.randn(RANK_LORA, HIDDEN),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)),
|
|
}
|
|
|
|
|
|
def sd_oft_kohya_proj():
|
|
"""Kohya-format OFT (oft_blocks) on a non-fused proj target."""
|
|
# OFT block-diagonal: out_dim = HIDDEN, num_blocks=8, block_size=16.
|
|
return {
|
|
'lora_unet_double_blocks_1_img_attn_proj.oft_blocks': torch.zeros(8, 16, 16),
|
|
'lora_unet_double_blocks_1_img_attn_proj.alpha': torch.tensor(1e-3),
|
|
}
|
|
|
|
|
|
def sd_oft_lycoris_proj():
|
|
"""LyCORIS-format OFT (oft_diag) — exercises the network_oft.py:58 NPE-fix path.
|
|
|
|
LyCORIS reads ``self.dim = oft_diag.shape[1]`` (the block_size) and computes
|
|
``block_size, num_blocks = factorization(out_dim, dim)``. ``factorization``
|
|
always returns ``(smaller, larger)``. For out_dim=HIDDEN=128 with dim=8, it
|
|
returns (8, 16) ⇒ block_size=8, num_blocks=16. The tensor must therefore
|
|
have shape ``(num_blocks, block_size, block_size) = (16, 8, 8)``.
|
|
"""
|
|
return {
|
|
'lora_unet_double_blocks_1_img_attn_proj.oft_diag': torch.zeros(16, 8, 8),
|
|
}
|
|
|
|
|
|
def sd_oft_kohya_qkv_skipped():
|
|
"""OFT on fused img_attn.qkv — must be skipped with a warning (no chunk class)."""
|
|
return {
|
|
'lora_unet_double_blocks_0_img_attn_qkv.oft_blocks': torch.zeros(8, 48, 48),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(1e-3),
|
|
}
|
|
|
|
|
|
def sd_boft_butterfly():
|
|
"""BOFT (butterfly-OFT) — reuses ``oft_blocks`` key with 4-D shape.
|
|
|
|
Mirrors the LyCORIS upstream save layout (boft.py weight_list:
|
|
``oft_blocks, rescale, alpha``; tensor shape per __init__:
|
|
``(boft_m, block_num, block_size, block_size)``).
|
|
|
|
For HIDDEN=128 with block_size=16, block_num=8:
|
|
- boft_m derives from sum-of-set-bits of (block_num-1) + 1 = bits(7)+1 = 4
|
|
- cascade indices: k = 2^i * r_b for i in 0..3, with r_b = 8
|
|
- all four stages divide cleanly into 128-row tensors
|
|
"""
|
|
return {
|
|
'lora_unet_double_blocks_1_img_attn_proj.oft_blocks':
|
|
torch.randn(4, 8, 16, 16) * 0.01,
|
|
'lora_unet_double_blocks_1_img_attn_proj.alpha':
|
|
torch.tensor(1e-3),
|
|
}
|
|
|
|
|
|
def sd_ia3_proj():
|
|
"""IA3 on a non-fused target. on_input=False ⇒ output-axis scaling."""
|
|
return {
|
|
'lora_unet_double_blocks_1_img_attn_proj.weight': torch.randn(HIDDEN),
|
|
'lora_unet_double_blocks_1_img_attn_proj.on_input': torch.tensor(0),
|
|
}
|
|
|
|
|
|
def sd_ia3_qkv_skipped():
|
|
"""IA3 on fused QKV — must be skipped with a warning."""
|
|
return {
|
|
'lora_unet_double_blocks_0_img_attn_qkv.weight': torch.randn(3 * QKV_OUT),
|
|
'lora_unet_double_blocks_0_img_attn_qkv.on_input': torch.tensor(0),
|
|
}
|
|
|
|
|
|
def sd_glora_proj():
|
|
"""GLoRA on a non-fused proj target.
|
|
|
|
Target host module is ``attn.to_out.0`` with weight shape (out, in)
|
|
= (HIDDEN, QKV_OUT). NetworkModuleGLora computes
|
|
``updown = (w2b @ w1b) + ((target @ w2a) @ w1a)``, which needs:
|
|
|
|
- w2b: (out, rank), w1b: (rank, in) ⇒ first term shape (out, in)
|
|
- w2a: (in, rank), w1a: (rank, in) ⇒ second term ((out,in) @ (in,rank)) @ (rank,in) = (out, in)
|
|
"""
|
|
out_dim = HIDDEN
|
|
in_dim = QKV_OUT
|
|
return {
|
|
'lora_unet_double_blocks_1_img_attn_proj.a1.weight': torch.randn(RANK_LORA, in_dim),
|
|
'lora_unet_double_blocks_1_img_attn_proj.a2.weight': torch.randn(in_dim, RANK_LORA),
|
|
'lora_unet_double_blocks_1_img_attn_proj.b1.weight': torch.randn(RANK_LORA, in_dim),
|
|
'lora_unet_double_blocks_1_img_attn_proj.b2.weight': torch.randn(out_dim, RANK_LORA),
|
|
'lora_unet_double_blocks_1_img_attn_proj.alpha': torch.tensor(float(RANK_LORA)),
|
|
}
|
|
|
|
|
|
def sd_norm_peft_attn_norm_q():
|
|
"""Norm targeting a real Flux2 RMSNorm (attn.norm_q)."""
|
|
return {
|
|
'transformer.transformer_blocks.0.attn.norm_q.w_norm': torch.randn(HEAD_DIM) * 0.01,
|
|
'transformer.transformer_blocks.0.attn.norm_q.b_norm': torch.randn(HEAD_DIM) * 0.01,
|
|
}
|
|
|
|
|
|
def sd_full_proj():
|
|
"""Full-rank delta on a proj target. diff matches host weight shape (out, in)."""
|
|
return {
|
|
'lora_unet_double_blocks_1_img_attn_proj.diff': torch.randn(HIDDEN, QKV_OUT) * 0.01,
|
|
'lora_unet_double_blocks_1_img_attn_proj.diff_b': torch.randn(HIDDEN) * 0.01,
|
|
}
|
|
|
|
|
|
# ============================================================
|
|
# Test plumbing helpers
|
|
# ============================================================
|
|
|
|
class TempLora:
|
|
"""Context-manager that materializes a synthetic state dict as a real
|
|
safetensors file and yields a NetworkOnDisk pointing at it.
|
|
|
|
Required because flux2_lora's loaders read via sd_models.read_state_dict,
|
|
which only accepts a filename. Cleans up the temp file on exit.
|
|
"""
|
|
|
|
def __init__(self, state_dict, name='test'):
|
|
self.state_dict = state_dict
|
|
self.name = name
|
|
self.path = None
|
|
|
|
def __enter__(self):
|
|
fd, self.path = tempfile.mkstemp(suffix='.safetensors', prefix=f'{self.name}_')
|
|
os.close(fd)
|
|
# safetensors requires contiguous tensors.
|
|
sd = {k: v.contiguous() for k, v in self.state_dict.items()}
|
|
safetensors.torch.save_file(sd, self.path)
|
|
return _MockNetworkOnDisk(self.name, self.path)
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
if self.path and os.path.exists(self.path):
|
|
os.unlink(self.path)
|
|
|
|
|
|
class _MockNetworkOnDisk:
|
|
"""Minimal NetworkOnDisk shim — loaders only touch .filename, .name, .shorthash, .sd_version."""
|
|
|
|
def __init__(self, name, filename):
|
|
self.name = name
|
|
self.filename = filename
|
|
self.shorthash = ''
|
|
self.sd_version = 'f2'
|
|
self.metadata = {}
|
|
|
|
def read_hash(self):
|
|
pass
|
|
|
|
|
|
def assert_shape(t: torch.Tensor, expected_shape, label=''):
|
|
actual = tuple(t.shape)
|
|
expected = tuple(expected_shape)
|
|
assert actual == expected, f'{label} shape {actual} != expected {expected}'
|
|
|
|
|
|
def assert_finite(t: torch.Tensor, label=''):
|
|
assert torch.isfinite(t).all(), f'{label} contains non-finite values'
|
|
|
|
|
|
def make_network_for_module(net_module: network.NetworkModule, te_mul: float = 1.0, unet_mul: float = 1.0):
|
|
"""Wire a NetworkModule into a Network with given multipliers so calc_updown
|
|
has a meaningful multiplier()/calc_scale() context."""
|
|
net_module.network.te_multiplier = te_mul
|
|
net_module.network.unet_multiplier = unet_mul
|
|
return net_module
|
|
|
|
|
|
# ============================================================
|
|
# Tests — parsing primitives
|
|
# ============================================================
|
|
|
|
CAT_PARSE = category('parse')
|
|
|
|
|
|
def test_parse_key_all_prefixes():
|
|
cases = [
|
|
('lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight',
|
|
F.LORA_SUFFIXES,
|
|
('lora_unet_', 'double_blocks_0_img_attn_qkv', 'lora_down.weight')),
|
|
('diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight',
|
|
F.LORA_SUFFIXES,
|
|
('diffusion_model.', 'double_blocks.0.img_attn.proj', 'lora_down.weight')),
|
|
('transformer.transformer_blocks.0.attn.to_q.lora_B.weight',
|
|
F.LORA_SUFFIXES,
|
|
('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_up.weight')),
|
|
('double_blocks.10.img_mlp.0.lora_A.weight',
|
|
F.LORA_SUFFIXES,
|
|
(F.BARE_DIFFUSERS_PREFIX_USED, 'double_blocks.10.img_mlp.0', 'lora_down.weight')),
|
|
('random.unrelated.key', F.LORA_SUFFIXES, None),
|
|
]
|
|
for key, suffixes, expected in cases:
|
|
got = F.parse_key(key, suffixes)
|
|
assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
|
|
return True
|
|
|
|
|
|
def test_resolve_targets_qkv_chunking():
|
|
from modules.lora.native_adapter import ChunkSpec
|
|
# Kohya double_blocks fused QKV → three chunks targeting Q/K/V.
|
|
targets = F.resolve_targets('lora_unet_', 'double_blocks_0_img_attn_qkv')
|
|
assert targets == [
|
|
('transformer_blocks.0.attn.to_q', ChunkSpec(idx=0, total=3)),
|
|
('transformer_blocks.0.attn.to_k', ChunkSpec(idx=1, total=3)),
|
|
('transformer_blocks.0.attn.to_v', ChunkSpec(idx=2, total=3)),
|
|
], f'kohya img_attn.qkv → {targets}'
|
|
|
|
targets = F.resolve_targets('lora_unet_', 'double_blocks_5_txt_attn_qkv')
|
|
assert targets == [
|
|
('transformer_blocks.5.attn.add_q_proj', ChunkSpec(idx=0, total=3)),
|
|
('transformer_blocks.5.attn.add_k_proj', ChunkSpec(idx=1, total=3)),
|
|
('transformer_blocks.5.attn.add_v_proj', ChunkSpec(idx=2, total=3)),
|
|
], f'kohya txt_attn.qkv → {targets}'
|
|
|
|
targets = F.resolve_targets('diffusion_model.', 'single_blocks.7.linear1')
|
|
assert targets == [('single_transformer_blocks.7.attn.to_qkv_mlp_proj', None)]
|
|
|
|
targets = F.native_adapter.resolve_group_targets(F.resolve_targets, 'transformer.', 'transformer_blocks.0.attn.to_q')
|
|
assert targets == [('transformer_blocks.0.attn.to_q', None)]
|
|
|
|
targets = F.resolve_targets('weird_prefix.', 'whatever')
|
|
assert targets == []
|
|
return True
|
|
|
|
|
|
def test_resolve_targets_extra_modules():
|
|
"""Every non-block F2_EXTRA_MAP target maps 1:1 (no chunk) in all three key forms."""
|
|
for bfl_base, diffusers_path in F.F2_EXTRA_MAP.items():
|
|
targets = F.resolve_targets('diffusion_model.', bfl_base)
|
|
assert targets == [(diffusers_path, None)], f'{bfl_base} -> {targets}'
|
|
targets = F.resolve_targets(F.BARE_DIFFUSERS_PREFIX_USED, bfl_base)
|
|
assert targets == [(diffusers_path, None)], f'bare {bfl_base} -> {targets}'
|
|
targets = F.resolve_targets('lora_unet_', bfl_base.replace('.', '_'))
|
|
assert targets == [(diffusers_path, None)], f'kohya {bfl_base} -> {targets}'
|
|
# guidance_in is a bare BFL path.
|
|
got = F.parse_key('guidance_in.in_layer.lora_A.weight', F.LORA_SUFFIXES)
|
|
assert got == (F.BARE_DIFFUSERS_PREFIX_USED, 'guidance_in.in_layer', 'lora_down.weight'), f'bare guidance_in parse -> {got}'
|
|
return True
|
|
|
|
|
|
def test_parse_key_peft_wrapper_unwrap():
|
|
"""base_model.model. wrapper is stripped before format detection.
|
|
|
|
Verifies the unwrap handles all four content cases:
|
|
- BFL keys under the wrapper (fal style)
|
|
- Diffusers paths under transformer. under the wrapper (HF DreamBooth style)
|
|
- Bare BFL under the wrapper
|
|
- Wrapper not present (passthrough)
|
|
"""
|
|
cases = [
|
|
('base_model.model.double_blocks.1.img_attn.proj.lora_A.weight',
|
|
F.LORA_SUFFIXES,
|
|
(F.BARE_DIFFUSERS_PREFIX_USED, 'double_blocks.1.img_attn.proj', 'lora_down.weight')),
|
|
('base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.weight',
|
|
F.LORA_SUFFIXES,
|
|
('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_down.weight')),
|
|
('base_model.model.lora_unet_double_blocks_0_img_attn_proj.lora_down.weight',
|
|
F.LORA_SUFFIXES,
|
|
('lora_unet_', 'double_blocks_0_img_attn_proj', 'lora_down.weight')),
|
|
('diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight', # no wrapper
|
|
F.LORA_SUFFIXES,
|
|
('diffusion_model.', 'double_blocks.0.img_attn.proj', 'lora_down.weight')),
|
|
]
|
|
for key, suffixes, expected in cases:
|
|
got = F.parse_key(key, suffixes)
|
|
assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
|
|
return True
|
|
|
|
|
|
def test_parse_key_lycoris_prefix():
|
|
"""lycoris_ prefix yields an underscored diffusers path; no _ -> . conversion."""
|
|
cases = [
|
|
('lycoris_transformer_blocks_0_attn_to_q.lokr_w1',
|
|
F.LOKR_SUFFIXES,
|
|
('lycoris_', 'transformer_blocks_0_attn_to_q', 'lokr_w1')),
|
|
('lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w2',
|
|
F.LOKR_SUFFIXES,
|
|
('lycoris_', 'transformer_blocks_0_attn_add_k_proj', 'lokr_w2')),
|
|
('lycoris_single_transformer_blocks_5_attn_to_qkv_mlp_proj.lokr_w1',
|
|
F.LOKR_SUFFIXES,
|
|
('lycoris_', 'single_transformer_blocks_5_attn_to_qkv_mlp_proj', 'lokr_w1')),
|
|
]
|
|
for key, suffixes, expected in cases:
|
|
got = F.parse_key(key, suffixes)
|
|
assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
|
|
|
|
# Resolution: lycoris_ is a universal passthrough handled upstream of the
|
|
# arch resolver, so the underscored path is returned verbatim (no chunk).
|
|
targets = F.native_adapter.resolve_group_targets(F.resolve_targets, 'lycoris_', 'transformer_blocks_0_attn_add_k_proj')
|
|
assert targets == [('transformer_blocks_0_attn_add_k_proj', None)], f'targets={targets}'
|
|
return True
|
|
|
|
|
|
def test_parse_key_bare_diffusers_and_peft_default():
|
|
"""Bare-diffusers paths + ``.lora_[AB].<adapter_name>.weight`` infix.
|
|
|
|
Covers ``Flux2Transformer2DModel.save_lora_adapter()`` output where the file
|
|
has no wrapping prefix and PEFT keeps the adapter slot name in the suffix.
|
|
"""
|
|
bd = F.BARE_DIFFUSERS_PREFIX_USED
|
|
cases = [
|
|
('single_transformer_blocks.0.attn.to_qkv_mlp_proj.lora_A.default.weight',
|
|
F.LORA_SUFFIXES,
|
|
(bd, 'single_transformer_blocks.0.attn.to_qkv_mlp_proj', 'lora_down.weight')),
|
|
('transformer_blocks.0.attn.to_q.lora_B.default.weight',
|
|
F.LORA_SUFFIXES,
|
|
(bd, 'transformer_blocks.0.attn.to_q', 'lora_up.weight')),
|
|
# Non-default adapter name also strips.
|
|
('transformer_blocks.7.attn.to_v.lora_A.style.weight',
|
|
F.LORA_SUFFIXES,
|
|
(bd, 'transformer_blocks.7.attn.to_v', 'lora_down.weight')),
|
|
# No adapter-name infix passes through unchanged.
|
|
('transformer_blocks.0.attn.to_q.lora_A.weight',
|
|
F.LORA_SUFFIXES,
|
|
(bd, 'transformer_blocks.0.attn.to_q', 'lora_down.weight')),
|
|
# Wrapped path with adapter-name infix: unwrap + strip both apply.
|
|
('base_model.model.transformer.transformer_blocks.0.attn.to_q.lora_A.default.weight',
|
|
F.LORA_SUFFIXES,
|
|
('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_down.weight')),
|
|
]
|
|
for key, suffixes, expected in cases:
|
|
got = F.parse_key(key, suffixes)
|
|
assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
|
|
|
|
# the shared resolver passes the bare-diffusers path through verbatim.
|
|
targets = F.native_adapter.resolve_group_targets(F.resolve_targets, bd, 'single_transformer_blocks.5.attn.to_out')
|
|
assert targets == [('single_transformer_blocks.5.attn.to_out', None)], f'targets={targets}'
|
|
return True
|
|
|
|
|
|
def test_marker_disambiguation():
|
|
# Each family's marker must reject other families' files.
|
|
pure_lora = {'lora_unet_x.lora_down.weight': torch.zeros(1, 1),
|
|
'lora_unet_x.lora_up.weight': torch.zeros(1, 1)}
|
|
assert F.has_marker(pure_lora, F.LORA_MARKERS)
|
|
assert not F.has_marker(pure_lora, F.LOKR_MARKERS)
|
|
assert not F.has_marker(pure_lora, F.LOHA_MARKERS)
|
|
assert not F.has_marker(pure_lora, F.OFT_MARKERS)
|
|
assert not F.has_marker(pure_lora, F.IA3_MARKERS)
|
|
assert not F.has_marker(pure_lora, F.GLORA_MARKERS)
|
|
assert not F.has_marker(pure_lora, F.NORM_MARKERS)
|
|
assert not F.has_marker(pure_lora, F.FULL_MARKERS)
|
|
|
|
ia3 = {'lora_unet_x.weight': torch.zeros(1), 'lora_unet_x.on_input': torch.tensor(1)}
|
|
assert F.has_marker(ia3, F.IA3_MARKERS)
|
|
# IA3 shares ".weight" but only on_input is the disambiguator.
|
|
return True
|
|
|
|
|
|
# ============================================================
|
|
# Tests — loaders (one per family)
|
|
# ============================================================
|
|
|
|
CAT_LOADER = category('loader')
|
|
|
|
|
|
def _load_via(try_fn, state_dict, name='test'):
|
|
install_mock_pipe()
|
|
with TempLora(state_dict, name=name) as nod:
|
|
return try_fn(name, nod, lora_scale=1.0)
|
|
|
|
|
|
def test_lora_kohya_fused_qkv_chunked():
|
|
net = _load_via(F.try_load_lora, sd_lora_kohya_qkv())
|
|
assert net is not None and len(net.modules) == 3, f'expected 3 chunked modules, got {net.modules if net else None}'
|
|
# Network keys correspond to the three split projections.
|
|
expected_keys = {
|
|
'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_keys, f'got {set(net.modules)}'
|
|
# Every module is a NetworkModuleLora; up-weight has been chunked along dim 0.
|
|
for nk, mod in net.modules.items():
|
|
assert isinstance(mod, network_lora.NetworkModuleLora), f'{nk}: type={type(mod).__name__}'
|
|
# The chunked up tensor has shape (QKV_OUT, RANK), not (3*QKV_OUT, RANK).
|
|
up_shape = tuple(mod.up_model.weight.shape)
|
|
assert up_shape == (QKV_OUT, RANK_LORA), f'{nk}: up shape {up_shape}'
|
|
return True
|
|
|
|
|
|
def test_lora_bfl_non_fused():
|
|
net = _load_via(F.try_load_lora, sd_lora_bfl_proj())
|
|
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
|
|
nk, mod = next(iter(net.modules.items()))
|
|
assert nk == 'lora_transformer_transformer_blocks_1_attn_to_out_0', f'{nk}'
|
|
assert isinstance(mod, network_lora.NetworkModuleLora)
|
|
return True
|
|
|
|
|
|
def test_lora_peft_format():
|
|
net = _load_via(F.try_load_lora, sd_lora_peft_to_q())
|
|
assert net is not None and len(net.modules) == 1
|
|
assert 'lora_transformer_transformer_blocks_0_attn_to_q' in net.modules
|
|
return True
|
|
|
|
|
|
def test_lora_onetrainer_diffusers_flat():
|
|
"""OneTrainer lora_transformer_ diffusers-flat keys load via the shared passthrough.
|
|
|
|
These keys are sdnext's own internal network_layer_mapping names, so they
|
|
bind with no rename or chunking. to_qkv_mlp_proj is a single diffusers
|
|
module here, so it loads as one passthrough target (not chunked)."""
|
|
sd = {
|
|
'lora_transformer_transformer_blocks_0_attn_to_q.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'lora_transformer_transformer_blocks_0_attn_to_q.lora_up.weight': torch.randn(QKV_OUT, RANK_LORA),
|
|
'lora_transformer_transformer_blocks_0_attn_to_q.alpha': torch.tensor(float(RANK_LORA)),
|
|
'lora_transformer_transformer_blocks_0_ff_linear_in.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'lora_transformer_transformer_blocks_0_ff_linear_in.lora_up.weight': torch.randn(MLP_OUT, RANK_LORA),
|
|
'lora_transformer_transformer_blocks_0_ff_linear_in.alpha': torch.tensor(float(RANK_LORA)),
|
|
'lora_transformer_single_transformer_blocks_0_attn_to_qkv_mlp_proj.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'lora_transformer_single_transformer_blocks_0_attn_to_qkv_mlp_proj.lora_up.weight': torch.randn(SINGLE_FUSED_OUT, RANK_LORA),
|
|
'lora_transformer_single_transformer_blocks_0_attn_to_qkv_mlp_proj.alpha': torch.tensor(float(RANK_LORA)),
|
|
}
|
|
net = _load_via(F.try_load_lora, sd)
|
|
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
|
|
assert set(net.modules) == {
|
|
'lora_transformer_transformer_blocks_0_attn_to_q',
|
|
'lora_transformer_transformer_blocks_0_ff_linear_in',
|
|
'lora_transformer_single_transformer_blocks_0_attn_to_qkv_mlp_proj',
|
|
}, f'got {set(net.modules)}'
|
|
return True
|
|
|
|
|
|
def test_lora_bare_bfl_format():
|
|
net = _load_via(F.try_load_lora, sd_lora_bare_bfl_mlp())
|
|
assert net is not None and len(net.modules) == 1
|
|
assert 'lora_transformer_transformer_blocks_0_ff_linear_in' in net.modules
|
|
return True
|
|
|
|
|
|
def test_lora_peft_saved_fal_style():
|
|
"""Real-world PEFT-saved LoRA with BFL paths under base_model.model. wrapper."""
|
|
net = _load_via(F.try_load_lora, sd_lora_peft_saved_fal_style())
|
|
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
|
|
# base_model.model.double_blocks.1.img_attn.proj → transformer_blocks.1.attn.to_out.0
|
|
assert 'lora_transformer_transformer_blocks_1_attn_to_out_0' in net.modules
|
|
return True
|
|
|
|
|
|
def test_lora_peft_saved_dreambooth_style():
|
|
"""PEFT-saved LoRA with diffusers paths under base_model.model. wrapper.
|
|
|
|
HF DreamBooth-style: the path under base_model.model. starts with
|
|
transformer., which the unwrap+reparse routes through the standard PEFT
|
|
branch. Verifies the unwrap handles nested format prefixes correctly.
|
|
"""
|
|
net = _load_via(F.try_load_lora, sd_lora_peft_saved_dreambooth_style())
|
|
assert net is not None and len(net.modules) == 1
|
|
assert 'lora_transformer_transformer_blocks_0_attn_to_q' in net.modules
|
|
return True
|
|
|
|
|
|
def test_lora_peft_saved_diffusers_style():
|
|
"""Bare-diffusers paths + PEFT ``.default.`` adapter-name infix.
|
|
|
|
Exercises the bare-diffusers prefix branch (``single_transformer_blocks.``,
|
|
``transformer_blocks.``) combined with the PEFT-saved ``.lora_A.default.weight``
|
|
suffix shape produced by ``Flux2Transformer2DModel.save_lora_adapter()``.
|
|
"""
|
|
net = _load_via(F.try_load_lora, sd_lora_peft_saved_diffusers_style())
|
|
assert net is not None and len(net.modules) == 2, f'expected 2 modules, got {net.modules if net else None}'
|
|
assert 'lora_transformer_single_transformer_blocks_0_attn_to_qkv_mlp_proj' in net.modules
|
|
assert 'lora_transformer_transformer_blocks_0_attn_to_q' in net.modules
|
|
return True
|
|
|
|
|
|
def test_lora_dora_threading():
|
|
net = _load_via(F.try_load_lora, sd_lora_with_dora_scale())
|
|
assert net is not None and len(net.modules) == 1
|
|
mod = next(iter(net.modules.values()))
|
|
assert mod.dora_scale is not None, 'dora_scale not threaded into NetworkModule'
|
|
# Apply path runs through finalize_updown which calls apply_weight_decompose.
|
|
target = torch.randn(QKV_OUT, HIDDEN)
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, (QKV_OUT, HIDDEN), 'DoRA updown')
|
|
assert_finite(updown, 'DoRA updown')
|
|
return True
|
|
|
|
|
|
def test_lora_dora_fused_qkv_sliced():
|
|
"""Per-output dora_scale is sliced with the up-weight chunk and applies cleanly."""
|
|
net = _load_via(F.try_load_lora, sd_lora_dora_fused_qkv())
|
|
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
|
|
for nk, mod in net.modules.items():
|
|
assert mod.dora_scale is not None and mod.dora_scale.shape[0] == QKV_OUT, \
|
|
f'{nk}: dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
|
|
updown, _ex_bias = mod.calc_updown(mod.sd_module.weight)
|
|
assert_shape(updown, mod.sd_module.weight.shape, label=nk)
|
|
assert_finite(updown, label=nk)
|
|
return True
|
|
|
|
|
|
def test_lokr_dora_fused_qkv_sliced():
|
|
"""Per-output dora_scale rides the kron chunk; each module sees its own rows."""
|
|
net = _load_via(F.try_load_lokr, sd_lokr_dora_fused_qkv())
|
|
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
|
|
for nk, mod in net.modules.items():
|
|
assert isinstance(mod, network_lokr.NetworkModuleLokrChunk), f'{nk}: type={type(mod).__name__}'
|
|
assert mod.dora_scale is not None and mod.dora_scale.shape[0] == QKV_OUT, \
|
|
f'{nk}: dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
|
|
updown, _ex_bias = mod.calc_updown(mod.sd_module.weight)
|
|
assert_shape(updown, mod.sd_module.weight.shape, label=nk)
|
|
assert_finite(updown, label=nk)
|
|
return True
|
|
|
|
|
|
def test_dora_per_input_fused_skipped():
|
|
"""wd_on_out=False dora_scale on a fused target cannot be split; the group is skipped."""
|
|
net = _load_via(F.try_load_lokr, sd_lokr_dora_fused_qkv(per_input=True))
|
|
assert net is None, f'expected no modules, got {net.modules if net else None}'
|
|
return True
|
|
|
|
|
|
def test_lora_diff_b_fused_qkv_sliced():
|
|
"""Per-output diff_b slices with the chunk and flows out as ex_bias."""
|
|
net = _load_via(F.try_load_lora, sd_lora_diff_b_fused_qkv())
|
|
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
|
|
for nk, mod in net.modules.items():
|
|
assert mod.ex_bias is not None and mod.ex_bias.shape[0] == QKV_OUT, \
|
|
f'{nk}: ex_bias shape {tuple(mod.ex_bias.shape) if mod.ex_bias is not None else None}'
|
|
_updown, ex_bias = mod.calc_updown(mod.sd_module.weight)
|
|
assert ex_bias is not None and ex_bias.shape[0] == QKV_OUT, f'{nk}: applied ex_bias {ex_bias.shape if ex_bias is not None else None}'
|
|
return True
|
|
|
|
|
|
def test_legacy_bias_fused_skipped():
|
|
"""The legacy weight-shaped bias key has no partition on fused targets; group skipped."""
|
|
sd = dict(sd_lora_dora_fused_qkv())
|
|
del sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale']
|
|
sd['lora_unet_double_blocks_0_img_attn_qkv.bias'] = torch.randn(3 * QKV_OUT, HIDDEN)
|
|
net = _load_via(F.try_load_lora, sd)
|
|
assert net is None, f'expected skip, got {net.modules if net else None}'
|
|
return True
|
|
|
|
|
|
def test_lora_aitk_magnitude_dora():
|
|
"""ai-toolkit `magnitude` converts onto the dora_scale path, (out, 1)-shaped.
|
|
|
|
Numeric check on a square layer: with the raw 1-D vector the apply-time
|
|
orientation detection would fall to per-input and renormalize columns;
|
|
the (out, 1) reshape pins per-output. Reference: (W+D) * m/||W+D||_row - W
|
|
at scale 1 (ai-toolkit DoRA saves no alpha key).
|
|
"""
|
|
torch.manual_seed(0)
|
|
w_base = torch.randn(HIDDEN, QKV_OUT)
|
|
sd, down, up, magnitude = sd_lora_aitk_magnitude_dora(w_base)
|
|
net = _load_via(F.try_load_lora, sd)
|
|
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
|
|
mod = next(iter(net.modules.values()))
|
|
assert mod.dora_scale is not None and tuple(mod.dora_scale.shape) == (HIDDEN, 1), \
|
|
f'dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
|
|
updown, _ex_bias = mod.calc_updown(w_base.clone())
|
|
merged = w_base + up @ down
|
|
norm = merged.reshape(HIDDEN, -1).norm(dim=1).reshape(HIDDEN, 1) + torch.finfo(w_base.dtype).eps
|
|
ref = merged * (magnitude.reshape(HIDDEN, 1) / norm) - w_base
|
|
rel = ((updown - ref).norm() / (ref.norm() + 1e-12)).item()
|
|
assert torch.allclose(updown, ref, rtol=1e-4, atol=1e-5), f'rel err {rel:.4f}'
|
|
return True
|
|
|
|
|
|
def test_lora_magnitude_fused_qkv_sliced():
|
|
"""A fused-qkv magnitude vector converts, reshapes and slices per chunk."""
|
|
sd = dict(sd_lora_dora_fused_qkv())
|
|
del sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale']
|
|
sd['lora_unet_double_blocks_0_img_attn_qkv.magnitude'] = torch.rand(3 * QKV_OUT) + 0.5
|
|
net = _load_via(F.try_load_lora, sd)
|
|
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
|
|
for nk, mod in net.modules.items():
|
|
assert mod.dora_scale is not None and tuple(mod.dora_scale.shape) == (QKV_OUT, 1), \
|
|
f'{nk}: dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
|
|
return True
|
|
|
|
|
|
def test_lora_peft_magnitude_vector():
|
|
"""PEFT/diffusers `lora_magnitude_vector` converts the same way."""
|
|
sd = {
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(QKV_OUT, RANK_LORA),
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_magnitude_vector': torch.rand(QKV_OUT) + 0.5,
|
|
}
|
|
net = _load_via(F.try_load_lora, sd)
|
|
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
|
|
mod = next(iter(net.modules.values()))
|
|
assert mod.dora_scale is not None and tuple(mod.dora_scale.shape) == (QKV_OUT, 1), \
|
|
f'dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}'
|
|
return True
|
|
|
|
|
|
def test_lora_sparse_bias_residual():
|
|
"""LyCORIS use_bias extraction triplet reconstructs into the dense-bias path.
|
|
|
|
Mirrors the extraction save exactly: residual sparsified COO with int16
|
|
indices, values from the weight-shaped remainder, alpha == rank (scale 1).
|
|
Reference: total delta = up @ down + residual.
|
|
"""
|
|
torch.manual_seed(0)
|
|
down = torch.randn(RANK_LORA, QKV_OUT)
|
|
up = torch.randn(HIDDEN, RANK_LORA)
|
|
residual = torch.randn(HIDDEN, QKV_OUT)
|
|
residual[torch.rand_like(residual) < 0.98] = 0.0 # extraction sparsity default
|
|
sparse = residual.to_sparse().coalesce()
|
|
sd = {
|
|
'diffusion_model.double_blocks.1.img_attn.proj.lora_A.weight': down,
|
|
'diffusion_model.double_blocks.1.img_attn.proj.lora_B.weight': up,
|
|
'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.bias_indices': sparse.indices().to(torch.int16),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.bias_values': sparse.values(),
|
|
'diffusion_model.double_blocks.1.img_attn.proj.bias_size': torch.tensor(residual.shape).to(torch.int16),
|
|
}
|
|
net = _load_via(F.try_load_lora, sd)
|
|
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
|
|
mod = next(iter(net.modules.values()))
|
|
assert mod.bias is not None and mod.bias.is_sparse, f'bias={type(mod.bias)}'
|
|
w_base = torch.randn(HIDDEN, QKV_OUT)
|
|
updown, _ex_bias = mod.calc_updown(w_base)
|
|
ref = up @ down + residual
|
|
rel = ((updown - ref).norm() / (ref.norm() + 1e-12)).item()
|
|
assert torch.allclose(updown, ref, rtol=1e-4, atol=1e-5), f'rel err {rel:.4f}'
|
|
return True
|
|
|
|
|
|
def test_lora_sparse_bias_fused_skipped():
|
|
"""The sparse residual triplet on a fused target is skipped like dense bias."""
|
|
sd = dict(sd_lora_kohya_qkv())
|
|
residual = torch.zeros(3 * QKV_OUT, HIDDEN)
|
|
residual[0, 0] = 1.0
|
|
sparse = residual.to_sparse().coalesce()
|
|
base = 'lora_unet_double_blocks_0_img_attn_qkv'
|
|
sd[f'{base}.bias_indices'] = sparse.indices().to(torch.int16)
|
|
sd[f'{base}.bias_values'] = sparse.values()
|
|
sd[f'{base}.bias_size'] = torch.tensor(residual.shape).to(torch.int16)
|
|
net = _load_via(F.try_load_lora, sd)
|
|
assert net is None, f'expected skip, got {net.modules if net else None}'
|
|
return True
|
|
|
|
|
|
def test_lokr_shape_mismatch_rejected():
|
|
"""Kron dims that disagree with the module are rejected at load, not at apply."""
|
|
sd = sd_lokr_bfl_proj()
|
|
# Double the w2 in-dim: kron becomes (QKV_OUT, 2*HIDDEN) vs module (HIDDEN, QKV_OUT).
|
|
sd['diffusion_model.double_blocks.1.img_attn.proj.lokr_w2'] = \
|
|
torch.randn(QKV_OUT // LOKR_W1_DIM, 2 * (HIDDEN // LOKR_W1_DIM))
|
|
net = _load_via(F.try_load_lokr, sd)
|
|
assert net is None, f'expected mismatch rejection, got {net.modules if net else None}'
|
|
return True
|
|
|
|
|
|
def test_lokr_fused_shape_mismatch_rejected():
|
|
"""Fused kron rows must cover total*out; a 2-of-3 sized fused delta is rejected."""
|
|
sd = dict(sd_lokr_kohya_qkv())
|
|
# Kron rows 2*QKV_OUT instead of 3*QKV_OUT: no valid 3-way equal chunk.
|
|
sd['lora_unet_double_blocks_0_img_attn_qkv.lokr_w2'] = \
|
|
torch.randn((2 * QKV_OUT) // 8, HIDDEN // 4)
|
|
net = _load_via(F.try_load_lokr, sd)
|
|
assert net is None, f'expected mismatch rejection, got {net.modules if net else None}'
|
|
return True
|
|
|
|
|
|
def test_lokr_bfl_non_fused():
|
|
net = _load_via(F.try_load_lokr, sd_lokr_bfl_proj())
|
|
assert net is not None and len(net.modules) == 1
|
|
nk, mod = next(iter(net.modules.items()))
|
|
assert isinstance(mod, network_lokr.NetworkModuleLokr) and not isinstance(mod, network_lokr.NetworkModuleLokrChunk), \
|
|
f'{nk}: type={type(mod).__name__}'
|
|
return True
|
|
|
|
|
|
def test_lokr_kohya_fused_qkv_chunked():
|
|
net = _load_via(F.try_load_lokr, sd_lokr_kohya_qkv())
|
|
assert net is not None and len(net.modules) == 3
|
|
for nk, mod in net.modules.items():
|
|
assert isinstance(mod, network_lokr.NetworkModuleLokrChunk), f'{nk}: type={type(mod).__name__}'
|
|
assert 0 <= mod.chunk_index < 3 and mod.num_chunks == 3
|
|
return True
|
|
|
|
|
|
def test_lokr_simpletuner_lycoris_format():
|
|
"""Real-world LyCORIS-standalone LoKR with lycoris_ prefix on diffusers paths.
|
|
|
|
Critical case: one of the targets is ``add_k_proj`` whose Linear name has
|
|
embedded underscores. Confirms the loader does NOT do a naive _ -> .
|
|
expansion on the base path; the underscored form must round-trip through
|
|
the network_layer_mapping lookup intact.
|
|
"""
|
|
net = _load_via(F.try_load_lokr, sd_lokr_simpletuner_lycoris_style())
|
|
assert net is not None and len(net.modules) == 2, f'got {net.modules if net else None}'
|
|
expected = {
|
|
'lora_transformer_transformer_blocks_0_attn_to_q',
|
|
'lora_transformer_transformer_blocks_0_attn_add_k_proj',
|
|
}
|
|
assert set(net.modules) == expected, f'got {set(net.modules)}'
|
|
for mod in net.modules.values():
|
|
assert isinstance(mod, network_lokr.NetworkModuleLokr)
|
|
return True
|
|
|
|
|
|
def test_lokr_bfl_extra_modules_full_matrix():
|
|
"""All non-block F2_EXTRA_MAP targets bind; placeholder alpha never scales.
|
|
|
|
Regression for the gap where resolve_targets only knew single_blocks /
|
|
double_blocks and silently dropped embedder / modulation / final-layer
|
|
groups as unmapped.
|
|
"""
|
|
net = _load_via(F.try_load_lokr, sd_lokr_bfl_extra_modules())
|
|
bound = len(net.modules) if net else 0
|
|
assert net is not None and bound == len(F.F2_EXTRA_MAP), \
|
|
f'bound {bound}/{len(F.F2_EXTRA_MAP)}: {sorted(net.modules) if net else []}'
|
|
for nk, mod in net.modules.items():
|
|
assert isinstance(mod, network_lokr.NetworkModuleLokr) and not isinstance(mod, network_lokr.NetworkModuleLokrChunk), \
|
|
f'{nk}: type={type(mod).__name__}'
|
|
# Full-matrix factors leave dim unset; the ~1e10 placeholder alpha must be ignored.
|
|
assert mod.dim is None and mod.calc_scale() == 1.0, f'{nk}: dim={mod.dim} scale={mod.calc_scale()}'
|
|
updown, _ex_bias = mod.calc_updown(mod.sd_module.weight)
|
|
assert_shape(updown, mod.sd_module.weight.shape, label=nk)
|
|
assert_finite(updown, label=nk)
|
|
return True
|
|
|
|
|
|
def test_loha_bfl_non_fused():
|
|
net = _load_via(F.try_load_loha, sd_loha_bfl_proj())
|
|
assert net is not None and len(net.modules) == 1
|
|
nk, mod = next(iter(net.modules.items()))
|
|
assert isinstance(mod, network_hada.NetworkModuleHada) and not isinstance(mod, network_hada.NetworkModuleHadaChunk), \
|
|
f'{nk}: type={type(mod).__name__}'
|
|
return True
|
|
|
|
|
|
def test_loha_kohya_fused_qkv_chunked():
|
|
net = _load_via(F.try_load_loha, sd_loha_kohya_qkv())
|
|
assert net is not None and len(net.modules) == 3
|
|
for nk, mod in net.modules.items():
|
|
assert isinstance(mod, network_hada.NetworkModuleHadaChunk), f'{nk}: type={type(mod).__name__}'
|
|
assert 0 <= mod.chunk_index < 3 and mod.num_chunks == 3
|
|
return True
|
|
|
|
|
|
def test_oft_kohya_non_fused():
|
|
net = _load_via(F.try_load_oft, sd_oft_kohya_proj())
|
|
assert net is not None and len(net.modules) == 1
|
|
_, mod = next(iter(net.modules.items()))
|
|
assert isinstance(mod, network_oft.NetworkModuleOFT)
|
|
assert mod.is_kohya is True and mod.constraint is not None
|
|
return True
|
|
|
|
|
|
def test_oft_lycoris_no_npe():
|
|
"""LyCORIS OFT (oft_diag) loads and calc_updown completes without
|
|
accessing self.constraint, which is None in the LyCORIS path.
|
|
"""
|
|
net = _load_via(F.try_load_oft, sd_oft_lycoris_proj())
|
|
assert net is not None and len(net.modules) == 1
|
|
_, mod = next(iter(net.modules.items()))
|
|
assert isinstance(mod, network_oft.NetworkModuleOFT)
|
|
assert mod.is_kohya is False and mod.constraint is None
|
|
target = mod.sd_module.weight # real (HIDDEN, QKV_OUT) tensor from the mock
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape), 'LyCORIS OFT updown')
|
|
assert_finite(updown, 'LyCORIS OFT updown')
|
|
return True
|
|
|
|
|
|
def test_oft_fused_qkv_skipped():
|
|
"""OFT on fused QKV must be dropped (with warning) — no chunk class exists."""
|
|
net = _load_via(F.try_load_oft, sd_oft_kohya_qkv_skipped())
|
|
# Loader returns None when no modules bound (skipped + nothing else).
|
|
assert net is None or len(net.modules) == 0
|
|
return True
|
|
|
|
|
|
def test_boft_loads_and_applies():
|
|
"""BOFT files (4-D oft_blocks) bind to NetworkModuleBOFT via the
|
|
try_load_oft dispatch. Verifies shape attrs derive correctly from
|
|
the 4-D tensor and the butterfly-cascade calc_updown returns a
|
|
finite delta of host weight shape.
|
|
"""
|
|
net = _load_via(F.try_load_oft, sd_boft_butterfly())
|
|
assert net is not None and len(net.modules) == 1, \
|
|
f'BOFT did not load (got {net.modules if net else None})'
|
|
mod = next(iter(net.modules.values()))
|
|
assert isinstance(mod, network_boft.NetworkModuleBOFT), \
|
|
f'expected NetworkModuleBOFT, got {type(mod).__name__}'
|
|
assert isinstance(mod, network_boft.NetworkModuleBOFT) and not isinstance(mod, network_oft.NetworkModuleOFT), \
|
|
'BOFT must not be a subclass of NetworkModuleOFT (separate algorithm)'
|
|
# BOFT-specific shape attributes resolved correctly
|
|
assert mod.boft_m == 4 and mod.block_num == 8 and mod.block_size == 16, \
|
|
f'shape attrs: boft_m={mod.boft_m} block_num={mod.block_num} block_size={mod.block_size}'
|
|
# End-to-end calc_updown via the butterfly cascade
|
|
target = mod.sd_module.weight # (HIDDEN=128, QKV_OUT=128) for to_out.0
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape), 'BOFT updown')
|
|
assert_finite(updown, 'BOFT updown')
|
|
return True
|
|
|
|
|
|
def test_ia3_basic_load():
|
|
net = _load_via(F.try_load_ia3, sd_ia3_proj())
|
|
assert net is not None and len(net.modules) == 1
|
|
_, mod = next(iter(net.modules.items()))
|
|
assert isinstance(mod, network_ia3.NetworkModuleIa3)
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape), 'IA3 updown')
|
|
return True
|
|
|
|
|
|
def test_ia3_fused_qkv_skipped():
|
|
net = _load_via(F.try_load_ia3, sd_ia3_qkv_skipped())
|
|
assert net is None or len(net.modules) == 0
|
|
return True
|
|
|
|
|
|
def test_glora_basic_load_and_alpha():
|
|
"""GLoRA loads, self.dim is populated from w1b.shape[0], and
|
|
calc_scale returns alpha/dim rather than the 1.0 fallback.
|
|
"""
|
|
state_dict = sd_glora_proj()
|
|
net = _load_via(F.try_load_glora, state_dict)
|
|
assert net is not None and len(net.modules) == 1
|
|
mod = next(iter(net.modules.values()))
|
|
assert isinstance(mod, network_glora.NetworkModuleGLora)
|
|
assert mod.dim is not None, 'GLoRA dim not set (regression)'
|
|
expected_dim = state_dict['lora_unet_double_blocks_1_img_attn_proj.b1.weight'].shape[0]
|
|
assert mod.dim == expected_dim, f'GLoRA dim {mod.dim} != {expected_dim}'
|
|
# alpha=RANK_LORA (8), dim=expected_dim (8) → calc_scale = 1.0; instructive but matches kohya convention.
|
|
expected_alpha = float(RANK_LORA)
|
|
assert mod.alpha == expected_alpha, f'alpha={mod.alpha}'
|
|
expected_scale = expected_alpha / expected_dim
|
|
actual_scale = mod.calc_scale()
|
|
assert abs(actual_scale - expected_scale) < 1e-6, f'calc_scale={actual_scale} expected {expected_scale} (alpha/dim)'
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape), 'GLoRA updown')
|
|
assert_finite(updown, 'GLoRA updown')
|
|
return True
|
|
|
|
|
|
def test_norm_loader_local_stamping():
|
|
"""Regression for the loader-local stamping behavior in try_load_norm.
|
|
|
|
lora_convert.assign_network_names_to_compvis_modules deliberately skips
|
|
setting module.network_layer_name on transformer norm modules. The Norm
|
|
loader must stamp them itself for any module it actually binds, so the
|
|
activation loop (which checks network_layer_name is not None) can apply
|
|
the delta. Verify both: pre-loader the norm module has no
|
|
network_layer_name; post-loader the bound target does.
|
|
"""
|
|
sd_model = install_mock_pipe()
|
|
norm_module = sd_model.pipe.transformer.transformer_blocks[0].attn.norm_q
|
|
assert not getattr(norm_module, 'network_layer_name', None), \
|
|
'mock norm should not have network_layer_name pre-load'
|
|
# First call assign_network_names manually to mirror what resolve_mapping does.
|
|
from modules.lora import lora_convert
|
|
lora_convert.assign_network_names_to_compvis_modules(sd_model.pipe)
|
|
# The carve-out skipped stamping for norm modules.
|
|
assert not getattr(norm_module, 'network_layer_name', None), \
|
|
'norm module should still be unstamped after assign_network_names (lora_convert.py:502 carve-out)'
|
|
# Now run the loader.
|
|
state_dict = sd_norm_peft_attn_norm_q()
|
|
with TempLora(state_dict, name='norm_test') as nod:
|
|
net = F.try_load_norm('norm_test', nod, lora_scale=1.0)
|
|
assert net is not None and len(net.modules) == 1, 'Norm loader produced no modules'
|
|
mod = next(iter(net.modules.values()))
|
|
assert isinstance(mod, network_norm.NetworkModuleNorm)
|
|
# Loader-local stamping: the bound module now has network_layer_name set.
|
|
assert getattr(norm_module, 'network_layer_name', None) == 'lora_transformer_transformer_blocks_0_attn_norm_q', \
|
|
f'Norm loader did not stamp network_layer_name (got {getattr(norm_module, "network_layer_name", None)!r})'
|
|
# calc_updown returns updown + ex_bias for w_norm + b_norm.
|
|
target = norm_module.weight
|
|
updown, ex_bias = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape), 'Norm updown')
|
|
assert ex_bias is not None, 'Norm ex_bias should be present when b_norm is set'
|
|
assert_shape(ex_bias, tuple(target.shape), 'Norm ex_bias')
|
|
return True
|
|
|
|
|
|
def test_full_basic_load():
|
|
state_dict = sd_full_proj()
|
|
net = _load_via(F.try_load_full, state_dict)
|
|
assert net is not None and len(net.modules) == 1
|
|
mod = next(iter(net.modules.values()))
|
|
assert isinstance(mod, network_full.NetworkModuleFull)
|
|
target = mod.sd_module.weight
|
|
updown, ex_bias = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape), 'Full updown')
|
|
assert ex_bias is not None and tuple(ex_bias.shape) == (HIDDEN,), f'Full ex_bias shape {ex_bias.shape if ex_bias is not None else None}'
|
|
return True
|
|
|
|
|
|
# ============================================================
|
|
# Tests — calc_updown sanity per family
|
|
# ============================================================
|
|
|
|
CAT_MATH = category('math')
|
|
|
|
|
|
def _instantiate_first(net):
|
|
"""Convenience: pull the only module out of a single-module Network."""
|
|
assert net is not None and len(net.modules) == 1
|
|
return next(iter(net.modules.values()))
|
|
|
|
|
|
def test_lora_calc_updown_shape():
|
|
mod = _instantiate_first(_load_via(F.try_load_lora, sd_lora_bfl_proj()))
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_lokr_calc_updown_shape():
|
|
mod = _instantiate_first(_load_via(F.try_load_lokr, sd_lokr_bfl_proj()))
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_lokr_chunk_calc_updown_shape():
|
|
"""LokrChunk produces a (QKV_OUT, HIDDEN) chunk from a (3*QKV_OUT, HIDDEN) full kron."""
|
|
net = _load_via(F.try_load_lokr, sd_lokr_kohya_qkv())
|
|
for mod in net.modules.values():
|
|
target = mod.sd_module.weight # split projection: shape (HIDDEN, QKV_OUT) for to_q etc.
|
|
updown, _ = mod.calc_updown(target)
|
|
# `target.shape` is (out=QKV_OUT, in=HIDDEN) for the diffusers Linear; the chunk class
|
|
# returns the chunk of kron(w1,w2) sliced along dim 0 (the fused-output axis).
|
|
assert_shape(updown, (QKV_OUT, HIDDEN))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_loha_calc_updown_shape():
|
|
mod = _instantiate_first(_load_via(F.try_load_loha, sd_loha_bfl_proj()))
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_loha_chunk_calc_updown_shape():
|
|
"""HadaChunk produces a (QKV_OUT, HIDDEN) chunk per Q/K/V."""
|
|
net = _load_via(F.try_load_loha, sd_loha_kohya_qkv())
|
|
for mod in net.modules.values():
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, (QKV_OUT, HIDDEN))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_oft_calc_updown_shape():
|
|
mod = _instantiate_first(_load_via(F.try_load_oft, sd_oft_kohya_proj()))
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_ia3_calc_updown_shape():
|
|
mod = _instantiate_first(_load_via(F.try_load_ia3, sd_ia3_proj()))
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
# IA3 returns target * w; output shape matches target.
|
|
assert_shape(updown, tuple(target.shape))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_glora_calc_updown_shape():
|
|
mod = _instantiate_first(_load_via(F.try_load_glora, sd_glora_proj()))
|
|
target = mod.sd_module.weight
|
|
updown, _ = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_full_calc_updown_shape():
|
|
mod = _instantiate_first(_load_via(F.try_load_full, sd_full_proj()))
|
|
target = mod.sd_module.weight
|
|
updown, ex_bias = mod.calc_updown(target)
|
|
assert_shape(updown, tuple(target.shape))
|
|
assert_shape(ex_bias, (HIDDEN,))
|
|
assert_finite(updown)
|
|
return True
|
|
|
|
|
|
def test_dora_ordering_matches_lycoris():
|
|
"""DoRA merge equals the LyCORIS forward reference at multiplier 1 and
|
|
lerps the full merged delta at other multipliers (ComfyUI semantics).
|
|
|
|
Reference mirrors lycoris locon forward + apply_weight_decompose with
|
|
wd_on_out=True: alpha/rank scales the diff BEFORE the row norms. Uses
|
|
alpha != rank (the kohya-style case the old decompose-before-scale
|
|
ordering got ~64% wrong).
|
|
"""
|
|
torch.manual_seed(0)
|
|
alpha = 1.0
|
|
w_base = torch.randn(QKV_OUT, HIDDEN)
|
|
down = torch.randn(RANK_LORA, HIDDEN)
|
|
up = torch.randn(QKV_OUT, RANK_LORA)
|
|
dora = w_base.reshape(QKV_OUT, -1).norm(dim=1, keepdim=True)
|
|
|
|
def reference_delta(mult):
|
|
diff = (up @ down) * (alpha / RANK_LORA)
|
|
merged = w_base + diff
|
|
norm = merged.reshape(QKV_OUT, -1).norm(dim=1).reshape(QKV_OUT, 1) + torch.finfo(w_base.dtype).eps
|
|
return (merged * (dora / norm) - w_base) * mult
|
|
|
|
sd = {
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': down,
|
|
'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': up,
|
|
'transformer.transformer_blocks.0.attn.to_q.alpha': torch.tensor(alpha),
|
|
'transformer.transformer_blocks.0.attn.to_q.dora_scale': dora,
|
|
}
|
|
for mult in (1.0, 0.5):
|
|
net = _load_via(F.try_load_lora, sd)
|
|
assert net is not None and len(net.modules) == 1
|
|
mod = make_network_for_module(next(iter(net.modules.values())), te_mul=mult, unet_mul=mult)
|
|
updown, _ex_bias = mod.calc_updown(w_base.clone())
|
|
ref = reference_delta(mult)
|
|
rel = ((updown - ref).norm() / (ref.norm() + 1e-12)).item()
|
|
assert torch.allclose(updown, ref, rtol=1e-4, atol=1e-5), f'mult={mult}: rel err {rel:.4f}'
|
|
return True
|
|
|
|
|
|
# ============================================================
|
|
# Tests — apply path / regressions in shared infra
|
|
# ============================================================
|
|
|
|
CAT_APPLY = category('apply')
|
|
|
|
|
|
def test_ex_bias_accumulation_two_norms():
|
|
"""Stacking two Norm adapters that both produce ex_bias accumulates
|
|
correctly: network_calc_weights sums their contributions instead of
|
|
raising on a tensor truthiness check.
|
|
"""
|
|
install_mock_pipe()
|
|
norm_module = shared.sd_model.pipe.transformer.transformer_blocks[0].attn.norm_q
|
|
|
|
# Build two Norm Networks targeting the same norm module.
|
|
state_dict_a = {
|
|
'transformer.transformer_blocks.0.attn.norm_q.w_norm': torch.full((HEAD_DIM,), 0.01),
|
|
'transformer.transformer_blocks.0.attn.norm_q.b_norm': torch.full((HEAD_DIM,), 0.02),
|
|
}
|
|
state_dict_b = {
|
|
'transformer.transformer_blocks.0.attn.norm_q.w_norm': torch.full((HEAD_DIM,), 0.03),
|
|
'transformer.transformer_blocks.0.attn.norm_q.b_norm': torch.full((HEAD_DIM,), 0.04),
|
|
}
|
|
nets = []
|
|
for label, sd in (('a', state_dict_a), ('b', state_dict_b)):
|
|
with TempLora(sd, name=f'norm_{label}') as nod:
|
|
n = F.try_load_norm(f'norm_{label}', nod, lora_scale=1.0)
|
|
n.te_multiplier = 1.0
|
|
n.unet_multiplier = 1.0
|
|
nets.append(n)
|
|
|
|
# Patch loaded_networks for the duration of this test.
|
|
saved_loaded = list(l_common.loaded_networks)
|
|
l_common.loaded_networks.clear()
|
|
l_common.loaded_networks.extend(nets)
|
|
try:
|
|
# network_calc_weights iterates loaded_networks and hits the
|
|
# `if batch_ex_bias:` line on the second adapter.
|
|
layer_name = norm_module.network_layer_name
|
|
# Bind layer_name in case prior test cleared it.
|
|
if not layer_name:
|
|
from modules.lora import lora_convert
|
|
lora_convert.assign_network_names_to_compvis_modules(shared.sd_model.pipe)
|
|
layer_name = 'lora_transformer_transformer_blocks_0_attn_norm_q'
|
|
norm_module.network_layer_name = layer_name
|
|
batch_updown, batch_ex_bias = lora_apply.network_calc_weights(norm_module, layer_name)
|
|
assert batch_updown is not None, 'batch_updown should be present (two w_norm contributions)'
|
|
assert batch_ex_bias is not None, 'batch_ex_bias should be present (two b_norm contributions)'
|
|
assert_shape(batch_ex_bias, (HEAD_DIM,))
|
|
# Two contributions of 0.02 + 0.04 = 0.06 (pre-multiplier).
|
|
# multiplier()=te_multiplier=1.0 for both; finalize_updown applies it to ex_bias.
|
|
expected_b = 0.02 + 0.04
|
|
assert torch.allclose(batch_ex_bias, torch.full_like(batch_ex_bias, expected_b), atol=1e-5), \
|
|
f'ex_bias accumulated to {batch_ex_bias[0].item():.6f}, expected {expected_b:.6f}'
|
|
finally:
|
|
l_common.loaded_networks.clear()
|
|
l_common.loaded_networks.extend(saved_loaded)
|
|
return True
|
|
|
|
|
|
# ============================================================
|
|
# Test runner
|
|
# ============================================================
|
|
|
|
|
|
def run_tests():
|
|
t0 = time.time()
|
|
|
|
log.warning('=== Parsing primitives ===')
|
|
for fn in [test_parse_key_all_prefixes, test_resolve_targets_qkv_chunking,
|
|
test_resolve_targets_extra_modules,
|
|
test_parse_key_peft_wrapper_unwrap, test_parse_key_lycoris_prefix,
|
|
test_parse_key_bare_diffusers_and_peft_default,
|
|
test_marker_disambiguation]:
|
|
run_test(CAT_PARSE, fn)
|
|
|
|
log.warning('=== Loaders ===')
|
|
for fn in [
|
|
test_lora_kohya_fused_qkv_chunked,
|
|
test_lora_bfl_non_fused,
|
|
test_lora_peft_format,
|
|
test_lora_onetrainer_diffusers_flat,
|
|
test_lora_bare_bfl_format,
|
|
test_lora_peft_saved_fal_style,
|
|
test_lora_peft_saved_dreambooth_style,
|
|
test_lora_peft_saved_diffusers_style,
|
|
test_lora_dora_threading,
|
|
test_lora_dora_fused_qkv_sliced,
|
|
test_lokr_dora_fused_qkv_sliced,
|
|
test_dora_per_input_fused_skipped,
|
|
test_lora_diff_b_fused_qkv_sliced,
|
|
test_legacy_bias_fused_skipped,
|
|
test_lora_aitk_magnitude_dora,
|
|
test_lora_magnitude_fused_qkv_sliced,
|
|
test_lora_peft_magnitude_vector,
|
|
test_lora_sparse_bias_residual,
|
|
test_lora_sparse_bias_fused_skipped,
|
|
test_lokr_shape_mismatch_rejected,
|
|
test_lokr_fused_shape_mismatch_rejected,
|
|
test_lokr_bfl_non_fused,
|
|
test_lokr_kohya_fused_qkv_chunked,
|
|
test_lokr_simpletuner_lycoris_format,
|
|
test_lokr_bfl_extra_modules_full_matrix,
|
|
test_loha_bfl_non_fused,
|
|
test_loha_kohya_fused_qkv_chunked,
|
|
test_oft_kohya_non_fused,
|
|
test_oft_lycoris_no_npe,
|
|
test_oft_fused_qkv_skipped,
|
|
test_boft_loads_and_applies,
|
|
test_ia3_basic_load,
|
|
test_ia3_fused_qkv_skipped,
|
|
test_glora_basic_load_and_alpha,
|
|
test_norm_loader_local_stamping,
|
|
test_full_basic_load,
|
|
]:
|
|
run_test(CAT_LOADER, fn)
|
|
|
|
log.warning('=== calc_updown shape sanity ===')
|
|
for fn in [
|
|
test_lora_calc_updown_shape,
|
|
test_lokr_calc_updown_shape,
|
|
test_lokr_chunk_calc_updown_shape,
|
|
test_loha_calc_updown_shape,
|
|
test_loha_chunk_calc_updown_shape,
|
|
test_oft_calc_updown_shape,
|
|
test_ia3_calc_updown_shape,
|
|
test_glora_calc_updown_shape,
|
|
test_full_calc_updown_shape,
|
|
test_dora_ordering_matches_lycoris,
|
|
]:
|
|
run_test(CAT_MATH, fn)
|
|
|
|
log.warning('=== Apply path / shared-infra regressions ===')
|
|
for fn in [test_ex_bias_accumulation_two_norms]:
|
|
run_test(CAT_APPLY, fn)
|
|
|
|
elapsed = time.time() - t0
|
|
|
|
log.warning('=== Results ===')
|
|
total_passed = 0
|
|
total_failed = 0
|
|
for cat, data in results.items():
|
|
total_passed += data['passed']
|
|
total_failed += data['failed']
|
|
status = 'PASS' if data['failed'] == 0 else 'FAIL'
|
|
log.info(f' {cat}: {data["passed"]} passed, {data["failed"]} failed [{status}]')
|
|
log.warning(f'Total: {total_passed} passed, {total_failed} failed in {elapsed:.2f}s')
|
|
if total_failed:
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
run_tests()
|