"""
Offline unit tests for Chroma native adapter loaders.
Covers the four native families currently supported by ``pipelines.chroma.chroma_lora``
(LoRA, LoKR, LoHA, OFT) plus DoRA threading via the universal
``NetworkModule.finalize_updown`` hook.
Tests build a mock Chroma-shaped transformer, write synthetic safetensors
files for each adapter format observed in the wild, and exercise the full
loader path including the Flux-to-diffusers rename table and the unique
single-block ``linear1`` unequal-chunk slicing.
Save formats are cross-referenced against real Chroma LoRAs:
- BFL / AI-toolkit (``diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight``):
e.g. ``Chroma - Lenovo UltraReal``
- kohya (``lora_unet_double_blocks_0_img_attn_proj.lora_down.weight``):
e.g. ``90s_anime_aesthetic_Chroma``
- PEFT (``transformer.transformer_blocks.0.attn.to_q.lora_down.weight``)
Chroma LoRAs are trained against the Flux block layout (``double_blocks``,
``single_blocks``) regardless of save format. The diffusers
``ChromaTransformer2DModel`` exposes split-attention modules at
``transformer_blocks.X.attn.{to_q,to_k,to_v,...}`` and
``single_transformer_blocks.X.{attn.*, proj_mlp, proj_out}``. The loader
path-rewrites Flux paths to diffusers names and handles two distinct
fused-weight layouts:
- **Equal chunks** (double_blocks img_attn.qkv / txt_attn.qkv at
``[HIDDEN, HIDDEN, HIDDEN]``): LoRA chunks at load via ``torch.chunk``;
LoKR defers via ``NetworkModuleLokrChunk``.
- **Unequal chunks** (single_blocks linear1 at
``[HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN]``): LoRA slices row ranges at load;
LoKR defers via ``NetworkModuleLokrSliceChunk``.
LoHA and OFT on fused targets are skipped with a warning (no slice variant).
The ``distilled_guidance_layer`` (Chroma's central modulation generator that
replaces Flux's per-block ``norm1.linear``) is a real module path that
``assign_network_names_to_compvis_modules`` registers, so LoRAs targeting it
pass through unchanged.
No running server required.
Usage:
python test/test-chroma-native-adapters.py
"""
import os
import sys
import tempfile
import time
import torch
import safetensors.torch
script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, script_dir)
os.chdir(script_dir)
os.environ['SD_INSTALL_QUIET'] = '1'
import modules.cmd_args
import installer
_orig_argv = sys.argv
sys.argv = [sys.argv[0]]
try:
modules.cmd_args.parse_args()
finally:
sys.argv = _orig_argv
installer.add_args(modules.cmd_args.parser)
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
from modules.errors import log
from modules import shared
from modules.lora import (
network, network_lora, network_lokr, network_hada, network_oft,
)
from modules.lora import lora_common as l_common
from pipelines.chroma import chroma_lora as C
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:
record(cat, False, name, f'exception: {e}')
import traceback
traceback.print_exc()
HIDDEN = 96
HEAD_DIM = 32
MLP_RATIO = 4
MLP_HIDDEN = HIDDEN * MLP_RATIO
QKV_FUSED_OUT = 3 * HIDDEN
LINEAR1_OUT = 3 * HIDDEN + MLP_HIDDEN
LINEAR2_IN = HIDDEN + MLP_HIDDEN
N_DOUBLE = 2
N_SINGLE = 2
class _Holder(torch.nn.Module):
"""Empty container module - we attach children dynamically."""
def build_double_block():
"""Mirror ``ChromaTransformerBlock``'s diffusers-side module layout.
Uses ``FluxAttention(added_kv_proj_dim=dim)`` so both img-side and
context-side QKV + output projections are present.
"""
block = _Holder()
block.attn = _Holder()
block.attn.to_q = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.to_k = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.to_v = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.to_out = torch.nn.ModuleList([
torch.nn.Linear(HIDDEN, HIDDEN, bias=True),
torch.nn.Dropout(0.0),
])
block.attn.add_q_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.add_k_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.add_v_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.to_add_out = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.norm_q = torch.nn.RMSNorm(HEAD_DIM)
block.attn.norm_k = torch.nn.RMSNorm(HEAD_DIM)
block.attn.norm_added_q = torch.nn.RMSNorm(HEAD_DIM)
block.attn.norm_added_k = torch.nn.RMSNorm(HEAD_DIM)
block.ff = _Holder()
block.ff.net = torch.nn.ModuleList()
proj_act = _Holder()
proj_act.proj = torch.nn.Linear(HIDDEN, MLP_HIDDEN, bias=True)
block.ff.net.append(proj_act)
block.ff.net.append(torch.nn.Dropout(0.0))
block.ff.net.append(torch.nn.Linear(MLP_HIDDEN, HIDDEN, bias=True))
block.ff_context = _Holder()
block.ff_context.net = torch.nn.ModuleList()
proj_act_ctx = _Holder()
proj_act_ctx.proj = torch.nn.Linear(HIDDEN, MLP_HIDDEN, bias=True)
block.ff_context.net.append(proj_act_ctx)
block.ff_context.net.append(torch.nn.Dropout(0.0))
block.ff_context.net.append(torch.nn.Linear(MLP_HIDDEN, HIDDEN, bias=True))
return block
def build_single_block():
"""Mirror ``ChromaSingleTransformerBlock`` - has proj_mlp + attn (pre_only) + proj_out."""
block = _Holder()
block.attn = _Holder()
block.attn.to_q = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.to_k = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.to_v = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
block.attn.norm_q = torch.nn.RMSNorm(HEAD_DIM)
block.attn.norm_k = torch.nn.RMSNorm(HEAD_DIM)
block.proj_mlp = torch.nn.Linear(HIDDEN, MLP_HIDDEN, bias=True)
block.proj_out = torch.nn.Linear(LINEAR2_IN, HIDDEN, bias=True)
return block
def build_mock_transformer():
"""Build a torch.nn.Module mimicking ``ChromaTransformer2DModel``."""
transformer = _Holder()
transformer.transformer_blocks = torch.nn.ModuleList([build_double_block() for _ in range(N_DOUBLE)])
transformer.single_transformer_blocks = torch.nn.ModuleList([build_single_block() for _ in range(N_SINGLE)])
transformer.distilled_guidance_layer = _Holder()
transformer.distilled_guidance_layer.in_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
transformer.distilled_guidance_layer.out_proj = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
transformer.distilled_guidance_layer.layers = torch.nn.ModuleList()
transformer.distilled_guidance_layer.norms = torch.nn.ModuleList()
for _ in range(2):
mlp = _Holder()
mlp.linear_1 = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
mlp.linear_2 = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
transformer.distilled_guidance_layer.layers.append(mlp)
transformer.distilled_guidance_layer.norms.append(torch.nn.RMSNorm(HIDDEN))
transformer.x_embedder = torch.nn.Linear(HIDDEN, HIDDEN)
transformer.context_embedder = torch.nn.Linear(HIDDEN, HIDDEN)
transformer.proj_out = torch.nn.Linear(HIDDEN, HIDDEN, bias=True)
return transformer
class _MockChromaPipeline:
"""Class name carries 'Chroma' so name-based model-type dispatch routes correctly."""
def __init__(self, transformer):
self.transformer = transformer
self.text_encoder = None
class _MockChromaSdModel:
"""Outer wrapper holding pipe + network_layer_mapping."""
def __init__(self, pipe):
self.pipe = pipe
self.network_layer_mapping = {}
self.embedding_db = None
self.__class__.__name__ = 'ChromaPipeline'
def install_mock_pipe():
"""Set shared.sd_model to a mock exposing a Chroma-shaped transformer.
Each test re-installs so any prior network_layer_name stamps don't leak.
Writes directly to model_data.sd_model to bypass the ModelData lock.
Also patches ``chroma_lora.QKV_DIMS`` and ``chroma_lora.LINEAR1_DIMS`` to
match the test mock's scaled-down ``HIDDEN`` / ``MLP_HIDDEN``. The module
hardcodes Chroma1-HD's 3072 / 12288, which mismatches small test tensors
and causes ``split_fused_lora_group``'s
``up.shape[0] != sum(dims)`` gate to reject every fused fixture.
"""
transformer = build_mock_transformer()
pipe = _MockChromaPipeline(transformer)
sd_model = _MockChromaSdModel(pipe)
from modules.modeldata import model_data
model_data.sd_model = sd_model
transformer.config = _ChromaConfig(num_layers=N_DOUBLE, num_single_layers=N_SINGLE)
C.QKV_DIMS = [HIDDEN, HIDDEN, HIDDEN]
C.LINEAR1_DIMS = [HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN]
return sd_model
class _ChromaConfig:
def __init__(self, num_layers, num_single_layers):
self.num_layers = num_layers
self.num_single_layers = num_single_layers
RANK_LORA = 8
LOKR_W1_DIM = 8
def sd_lora_bfl_img_attn_proj():
"""BFL LoRA on double-block img_attn.proj.
BFL path maps to diffusers ``transformer_blocks.0.attn.to_out.0`` via
``DOUBLE_RENAME_TEMPLATES``.
"""
return {
'diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.double_blocks.0.img_attn.proj.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
'diffusion_model.double_blocks.0.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
}
def sd_lora_bfl_img_attn_qkv_fused():
"""BFL LoRA on fused img_attn.qkv. Loader splits up-weight along dim 0 into Q/K/V."""
return {
'diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight': torch.randn(QKV_FUSED_OUT, RANK_LORA),
'diffusion_model.double_blocks.0.img_attn.qkv.alpha': torch.tensor(float(RANK_LORA)),
}
def sd_lora_bfl_txt_attn_qkv_fused():
"""BFL LoRA on fused txt_attn.qkv. Loader emits 3 chunks to add_{q,k,v}_proj."""
return {
'diffusion_model.double_blocks.0.txt_attn.qkv.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.double_blocks.0.txt_attn.qkv.lora_B.weight': torch.randn(QKV_FUSED_OUT, RANK_LORA),
}
def sd_lora_bfl_img_mlp():
"""BFL LoRA on double-block img_mlp.0 and img_mlp.2.
img_mlp.0 -> ff.net.0.proj, img_mlp.2 -> ff.net.2.
"""
return {
'diffusion_model.double_blocks.1.img_mlp.0.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.double_blocks.1.img_mlp.0.lora_B.weight': torch.randn(MLP_HIDDEN, RANK_LORA),
'diffusion_model.double_blocks.1.img_mlp.2.lora_A.weight': torch.randn(RANK_LORA, MLP_HIDDEN),
'diffusion_model.double_blocks.1.img_mlp.2.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
}
def sd_lora_bfl_txt_mlp():
"""BFL LoRA on double-block txt_mlp.0 and txt_mlp.2 - context side."""
return {
'diffusion_model.double_blocks.0.txt_mlp.0.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.double_blocks.0.txt_mlp.0.lora_B.weight': torch.randn(MLP_HIDDEN, RANK_LORA),
}
def sd_lora_bfl_single_linear1_unequal():
"""BFL LoRA on single-block linear1.
linear1 fuses Q/K/V/proj_mlp at unequal dims [HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN].
Loader emits 4 targets with unequal row-range chunks.
"""
return {
'diffusion_model.single_blocks.0.linear1.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.single_blocks.0.linear1.lora_B.weight': torch.randn(LINEAR1_OUT, RANK_LORA),
}
def sd_lora_bfl_single_linear2():
"""BFL LoRA on single-block linear2 (-> single_transformer_blocks.X.proj_out)."""
return {
'diffusion_model.single_blocks.0.linear2.lora_A.weight': torch.randn(RANK_LORA, LINEAR2_IN),
'diffusion_model.single_blocks.0.linear2.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
}
def sd_lora_kohya_img_attn_proj():
"""Kohya flat-underscore LoRA on img_attn.proj. Mirrors 90s_anime_aesthetic_Chroma."""
return {
'lora_unet_double_blocks_0_img_attn_proj.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
'lora_unet_double_blocks_0_img_attn_proj.lora_up.weight': torch.randn(HIDDEN, RANK_LORA),
'lora_unet_double_blocks_0_img_attn_proj.alpha': torch.tensor(float(RANK_LORA)),
}
def sd_lora_kohya_img_attn_qkv_fused():
"""Kohya LoRA on fused img_attn.qkv."""
return {
'lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
'lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight': torch.randn(QKV_FUSED_OUT, RANK_LORA),
'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)),
}
def sd_lora_peft_to_q():
"""PEFT-format LoRA targeting a split diffusers path (no rename, no chunking)."""
return {
'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
}
def sd_lora_onetrainer_diffusers_flat():
"""OneTrainer LoRA: ``lora_transformer_`` + underscore-flat diffusers path.
QKV is pre-split (no fused chunks) and each key is byte-identical to
sdnext's internal network_layer_mapping entry, so resolve_targets passes
the base through unchanged. Mirrors a real OneTrainer save.
"""
return {
'lora_transformer_single_transformer_blocks_0_attn_to_q.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
'lora_transformer_single_transformer_blocks_0_attn_to_q.lora_up.weight': torch.randn(HIDDEN, RANK_LORA),
'lora_transformer_single_transformer_blocks_0_attn_to_q.alpha': torch.tensor(float(RANK_LORA)),
'lora_transformer_transformer_blocks_0_attn_to_out_0.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
'lora_transformer_transformer_blocks_0_attn_to_out_0.lora_up.weight': torch.randn(HIDDEN, RANK_LORA),
'lora_transformer_transformer_blocks_0_attn_to_out_0.alpha': torch.tensor(float(RANK_LORA)),
'lora_transformer_transformer_blocks_0_attn_add_k_proj.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
'lora_transformer_transformer_blocks_0_attn_add_k_proj.lora_up.weight': torch.randn(HIDDEN, RANK_LORA),
'lora_transformer_transformer_blocks_0_attn_add_k_proj.alpha': torch.tensor(float(RANK_LORA)),
'lora_transformer_transformer_blocks_0_ff_net_0_proj.lora_down.weight': torch.randn(RANK_LORA, HIDDEN),
'lora_transformer_transformer_blocks_0_ff_net_0_proj.lora_up.weight': torch.randn(MLP_HIDDEN, RANK_LORA),
'lora_transformer_transformer_blocks_0_ff_net_0_proj.alpha': torch.tensor(float(RANK_LORA)),
'lora_transformer_transformer_blocks_0_ff_net_2.lora_down.weight': torch.randn(RANK_LORA, MLP_HIDDEN),
'lora_transformer_transformer_blocks_0_ff_net_2.lora_up.weight': torch.randn(HIDDEN, RANK_LORA),
'lora_transformer_transformer_blocks_0_ff_net_2.alpha': torch.tensor(float(RANK_LORA)),
}
def sd_lora_distilled_guidance():
"""LoRA targeting Chroma's distilled_guidance_layer (passes through unchanged)."""
return {
'diffusion_model.distilled_guidance_layer.in_proj.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.distilled_guidance_layer.in_proj.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
}
def sd_lora_with_dora_scale():
"""LoRA with dora_scale companion to exercise DoRA threading."""
return {
'transformer.transformer_blocks.0.attn.to_v.lora_A.weight': torch.randn(RANK_LORA, HIDDEN),
'transformer.transformer_blocks.0.attn.to_v.lora_B.weight': torch.randn(HIDDEN, RANK_LORA),
'transformer.transformer_blocks.0.attn.to_v.dora_scale': torch.randn(HIDDEN),
}
def sd_lokr_bfl_img_attn_proj():
"""BFL LoKR on a non-fused proj target. Loader uses NetworkModuleLokr (no chunk)."""
return {
'diffusion_model.double_blocks.0.img_attn.proj.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
'diffusion_model.double_blocks.0.img_attn.proj.lokr_w2': torch.randn(HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
'diffusion_model.double_blocks.0.img_attn.proj.alpha': torch.tensor(float(LOKR_W1_DIM)),
}
def sd_lokr_bfl_img_attn_qkv_equal_chunks():
"""BFL LoKR on fused img_attn.qkv (equal chunks).
Loader emits 3 NetworkModuleLokrSliceChunk via row ranges [0:HIDDEN], [HIDDEN:2*HIDDEN], [2*HIDDEN:3*HIDDEN].
(Chroma's implementation slices even equal-chunks via row ranges since the same logic handles both.)
"""
return {
'diffusion_model.double_blocks.0.img_attn.qkv.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
'diffusion_model.double_blocks.0.img_attn.qkv.lokr_w2': torch.randn(QKV_FUSED_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
'diffusion_model.double_blocks.0.img_attn.qkv.alpha': torch.tensor(float(LOKR_W1_DIM)),
}
def sd_lokr_bfl_single_linear1_unequal():
"""BFL LoKR on fused single-block linear1 (UNEQUAL chunks).
Loader emits 4 NetworkModuleLokrSliceChunk with row ranges matching
[HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN] partitions.
"""
return {
'diffusion_model.single_blocks.0.linear1.lokr_w1': torch.randn(LOKR_W1_DIM, LOKR_W1_DIM),
'diffusion_model.single_blocks.0.linear1.lokr_w2': torch.randn(LINEAR1_OUT // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM),
'diffusion_model.single_blocks.0.linear1.alpha': torch.tensor(float(LOKR_W1_DIM)),
}
def sd_loha_bfl_img_attn_proj():
"""LoHA on a non-fused target binds via NetworkModuleHada."""
return {
'diffusion_model.double_blocks.1.img_attn.proj.hada_w1_a': torch.randn(HIDDEN, RANK_LORA),
'diffusion_model.double_blocks.1.img_attn.proj.hada_w1_b': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.double_blocks.1.img_attn.proj.hada_w2_a': torch.randn(HIDDEN, RANK_LORA),
'diffusion_model.double_blocks.1.img_attn.proj.hada_w2_b': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.double_blocks.1.img_attn.proj.alpha': torch.tensor(float(RANK_LORA)),
}
def sd_loha_bfl_img_attn_qkv_skipped():
"""LoHA on fused img_attn.qkv is dropped by the loader (no slice variant for LoHA)."""
return {
'diffusion_model.double_blocks.0.img_attn.qkv.hada_w1_a': torch.randn(QKV_FUSED_OUT, RANK_LORA),
'diffusion_model.double_blocks.0.img_attn.qkv.hada_w1_b': torch.randn(RANK_LORA, HIDDEN),
'diffusion_model.double_blocks.0.img_attn.qkv.hada_w2_a': torch.randn(QKV_FUSED_OUT, RANK_LORA),
'diffusion_model.double_blocks.0.img_attn.qkv.hada_w2_b': torch.randn(RANK_LORA, HIDDEN),
}
def sd_oft_bfl_img_attn_proj():
"""OFT (LyCORIS oft_diag form) on non-fused target."""
num_blocks = 4
block_size = HIDDEN // num_blocks
return {
'diffusion_model.double_blocks.0.img_attn.proj.oft_blocks': torch.randn(num_blocks, block_size, block_size) * 0.01,
'diffusion_model.double_blocks.0.img_attn.proj.oft_diag': torch.ones(num_blocks, block_size),
'diffusion_model.double_blocks.0.img_attn.proj.alpha': torch.tensor(0.001),
}
def sd_oft_bfl_img_attn_qkv_skipped():
"""OFT on fused img_attn.qkv - dropped by the loader (OFT structure tied to out_features)."""
num_blocks = 4
block_size = QKV_FUSED_OUT // num_blocks
return {
'diffusion_model.double_blocks.0.img_attn.qkv.oft_blocks': torch.randn(num_blocks, block_size, block_size) * 0.01,
}
class TempLora:
"""Context manager: writes a state dict to a temp safetensors file."""
def __init__(self, state_dict, name='test'):
self.state_dict = state_dict
self.name = name
self.path = None
def __enter__(self):
sd = {k: v.contiguous() if isinstance(v, torch.Tensor) else v for k, v in self.state_dict.items()}
fd, self.path = tempfile.mkstemp(suffix='.safetensors', prefix=f'{self.name}_')
os.close(fd)
safetensors.torch.save_file(sd, self.path)
return _MockNetworkOnDisk(self.path, self.name)
def __exit__(self, exc_type, exc_val, exc_tb):
if self.path and os.path.exists(self.path):
os.unlink(self.path)
class _MockNetworkOnDisk:
def __init__(self, filename, name):
self.filename = filename
self.name = name
self.shorthash = ''
self.sd_version = 'unknown'
def assert_shape(t: torch.Tensor, expected_shape, label=''):
actual = tuple(t.shape)
assert actual == tuple(expected_shape), f'{label}: shape {actual}, expected {tuple(expected_shape)}'
def make_network_for_module(net_module: network.NetworkModule, te_mul: float = 1.0, unet_mul: float = 1.0):
net_module.network.te_multiplier = te_mul
net_module.network.unet_multiplier = unet_mul
return net_module
CAT_PARSE = category('parse')
def test_parse_key_all_prefixes():
"""parse_key returns (prefix_used, base, suffix). Rename to diffusers happens
in resolve_targets, not parse_key."""
cases = [
('diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight',
C.LORA_SUFFIXES,
('diffusion_model.', 'double_blocks.0.img_attn.proj', 'lora_down.weight')),
('transformer.transformer_blocks.0.attn.to_q.lora_B.weight',
C.LORA_SUFFIXES,
('transformer.', 'transformer_blocks.0.attn.to_q', 'lora_up.weight')),
('lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight',
C.LORA_SUFFIXES,
('lora_unet_', 'double_blocks_0_img_attn_qkv', 'lora_down.weight')),
('lora_transformer_single_transformer_blocks_0_attn_to_q.lora_down.weight',
C.LORA_SUFFIXES,
('lora_transformer_', 'single_transformer_blocks_0_attn_to_q', 'lora_down.weight')),
('double_blocks.0.img_attn.proj.lora_A.weight',
C.LORA_SUFFIXES,
(C.BARE_DIFFUSERS_PREFIX_USED, 'double_blocks.0.img_attn.proj', 'lora_down.weight')),
('random.unrelated.key', C.LORA_SUFFIXES, None),
]
for key, suffixes, expected in cases:
got = C.parse_key(key, suffixes)
assert got == expected, f'parse_key({key!r}) = {got}, expected {expected}'
return True
def test_marker_disambiguation():
"""Each family's markers reject other families' files."""
pure_lora = {
'lora_unet_double_blocks_0_img_attn_proj.lora_down.weight': torch.zeros(1, 1),
'lora_unet_double_blocks_0_img_attn_proj.lora_up.weight': torch.zeros(1, 1),
}
assert C.has_marker(pure_lora, C.LORA_MARKERS)
assert not C.has_marker(pure_lora, C.LOKR_MARKERS)
assert not C.has_marker(pure_lora, C.LOHA_MARKERS)
assert not C.has_marker(pure_lora, C.OFT_MARKERS)
pure_lokr = {
'diffusion_model.double_blocks.0.img_attn.proj.lokr_w1': torch.zeros(1, 1),
'diffusion_model.double_blocks.0.img_attn.proj.lokr_w2': torch.zeros(1, 1),
}
assert C.has_marker(pure_lokr, C.LOKR_MARKERS)
assert not C.has_marker(pure_lokr, C.LORA_MARKERS)
return True
def test_resolve_targets_static_renames():
"""resolve_targets produces the documented Flux-to-diffusers remappings
for non-fused targets in both kohya and BFL forms.
"""
cases = [
(('lora_unet_', 'double_blocks_0_img_attn_proj'), 'transformer_blocks.0.attn.to_out.0'),
(('lora_unet_', 'double_blocks_0_txt_attn_proj'), 'transformer_blocks.0.attn.to_add_out'),
(('lora_unet_', 'double_blocks_1_img_mlp_0'), 'transformer_blocks.1.ff.net.0.proj'),
(('lora_unet_', 'double_blocks_1_img_mlp_2'), 'transformer_blocks.1.ff.net.2'),
(('lora_unet_', 'double_blocks_0_txt_mlp_0'), 'transformer_blocks.0.ff_context.net.0.proj'),
(('lora_unet_', 'single_blocks_0_linear2'), 'single_transformer_blocks.0.proj_out'),
(('diffusion_model.', 'double_blocks.0.img_attn.proj'), 'transformer_blocks.0.attn.to_out.0'),
(('diffusion_model.', 'single_blocks.0.linear2'), 'single_transformer_blocks.0.proj_out'),
]
for (prefix, base), expected_path in cases:
targets = C.resolve_targets(prefix, base)
assert len(targets) == 1, f'({prefix}, {base}) -> {targets}'
path, chunk = targets[0]
assert path == expected_path and chunk is None, f'({prefix}, {base}) -> {targets}'
return True
def test_resolve_targets_extra_and_guidance():
"""Non-block extra-map renames and guidance-layer MLP leaf renames, all key forms."""
for bfl_base, diffusers_path in C.CHROMA_EXTRA_MAP.items():
for prefix, base in [
('diffusion_model.', bfl_base),
(C.BARE_DIFFUSERS_PREFIX_USED, bfl_base),
('lora_unet_', bfl_base.replace('.', '_')),
]:
targets = C.resolve_targets(prefix, base)
assert targets == [(diffusers_path, None)], f'({prefix}, {base}) -> {targets}'
cases = [
(('diffusion_model.', 'distilled_guidance_layer.layers.0.in_layer'), 'distilled_guidance_layer.layers.0.linear_1'),
((C.BARE_DIFFUSERS_PREFIX_USED, 'distilled_guidance_layer.layers.1.out_layer'), 'distilled_guidance_layer.layers.1.linear_2'),
(('lora_unet_', 'distilled_guidance_layer_layers_0_in_layer'), 'distilled_guidance_layer_layers_0_linear_1'),
(('lora_unet_', 'distilled_guidance_layer_layers_1_out_layer'), 'distilled_guidance_layer_layers_1_linear_2'),
(('diffusion_model.', 'distilled_guidance_layer.in_proj'), 'distilled_guidance_layer.in_proj'),
((C.BARE_DIFFUSERS_PREFIX_USED, 'distilled_guidance_layer.layers.0.linear_1'), 'distilled_guidance_layer.layers.0.linear_1'),
]
for (prefix, base), expected in cases:
targets = C.resolve_targets(prefix, base)
assert targets == [(expected, None)], f'({prefix}, {base}) -> {targets}'
return True
def test_resolve_targets_onetrainer_passthrough():
"""The ``lora_transformer_`` passthrough lives in the shared resolver.
``lora_transformer_`` is sdnext's own internal transformer namespace.
``native_adapter.resolve_group_targets`` resolves it to an identity
passthrough for any arch; chroma's ``resolve_targets`` owns the Flux-layout,
bare, and kohya prefixes and returns nothing for it.
"""
na = C.native_adapter
assert C.resolve_targets('lora_transformer_', 'transformer_blocks_0_attn_to_q') == []
for base in [
'single_transformer_blocks_0_attn_to_q',
'transformer_blocks_0_attn_to_out_0',
'transformer_blocks_0_attn_add_k_proj',
'transformer_blocks_0_ff_net_0_proj',
'transformer_blocks_0_ff_net_2',
]:
targets = na.resolve_group_targets(C.resolve_targets, 'lora_transformer_', base)
assert targets == [(base, None)], f'{base} -> {targets}'
return True
CAT_LOADER = category('loader')
def _load_via(try_fn, state_dict, name='test'):
install_mock_pipe()
with TempLora(state_dict, name=name) as nod:
return try_fn(name, nod, lora_scale=1.0)
def test_lora_bfl_img_attn_proj():
"""BFL LoRA on img_attn.proj renames to attn.to_out.0."""
net = _load_via(C.try_load_lora, sd_lora_bfl_img_attn_proj())
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
assert 'lora_transformer_transformer_blocks_0_attn_to_out_0' in net.modules
return True
def test_lora_bfl_img_attn_qkv_chunked():
"""BFL LoRA on fused img_attn.qkv emits 3 chunks targeting to_q/to_k/to_v."""
net = _load_via(C.try_load_lora, sd_lora_bfl_img_attn_qkv_fused())
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
expected = {
'lora_transformer_transformer_blocks_0_attn_to_q',
'lora_transformer_transformer_blocks_0_attn_to_k',
'lora_transformer_transformer_blocks_0_attn_to_v',
}
assert set(net.modules) == expected
for nk, mod in net.modules.items():
assert_shape(mod.up_model.weight, (HIDDEN, RANK_LORA), label=nk)
return True
def test_lora_bfl_txt_attn_qkv_chunked():
"""BFL LoRA on fused txt_attn.qkv emits 3 chunks targeting add_q/k/v_proj (context side)."""
net = _load_via(C.try_load_lora, sd_lora_bfl_txt_attn_qkv_fused())
assert net is not None and len(net.modules) == 3
expected = {
'lora_transformer_transformer_blocks_0_attn_add_q_proj',
'lora_transformer_transformer_blocks_0_attn_add_k_proj',
'lora_transformer_transformer_blocks_0_attn_add_v_proj',
}
assert set(net.modules) == expected
return True
def test_lora_bfl_img_mlp():
"""img_mlp.0 -> ff.net.0.proj, img_mlp.2 -> ff.net.2."""
net = _load_via(C.try_load_lora, sd_lora_bfl_img_mlp())
assert net is not None and len(net.modules) == 2
assert 'lora_transformer_transformer_blocks_1_ff_net_0_proj' in net.modules
assert 'lora_transformer_transformer_blocks_1_ff_net_2' in net.modules
return True
def test_lora_bfl_txt_mlp():
"""txt_mlp.0 -> ff_context.net.0.proj."""
net = _load_via(C.try_load_lora, sd_lora_bfl_txt_mlp())
assert net is not None and len(net.modules) == 1
assert 'lora_transformer_transformer_blocks_0_ff_context_net_0_proj' in net.modules
return True
def test_lora_bfl_single_linear1_unequal_chunks():
"""BFL LoRA on single linear1 emits 4 targets with UNEQUAL row ranges.
Partitions: [HIDDEN, HIDDEN, HIDDEN, MLP_HIDDEN] -> to_q, to_k, to_v, proj_mlp.
The first three chunks have (HIDDEN, RANK) up-shape; the fourth has (MLP_HIDDEN, RANK).
"""
net = _load_via(C.try_load_lora, sd_lora_bfl_single_linear1_unequal())
assert net is not None and len(net.modules) == 4, f'got {net.modules if net else None}'
expected = {
'lora_transformer_single_transformer_blocks_0_attn_to_q',
'lora_transformer_single_transformer_blocks_0_attn_to_k',
'lora_transformer_single_transformer_blocks_0_attn_to_v',
'lora_transformer_single_transformer_blocks_0_proj_mlp',
}
assert set(net.modules) == expected
for nk, mod in net.modules.items():
if nk.endswith('proj_mlp'):
assert_shape(mod.up_model.weight, (MLP_HIDDEN, RANK_LORA), label=nk)
else:
assert_shape(mod.up_model.weight, (HIDDEN, RANK_LORA), label=nk)
return True
def test_lora_bfl_single_linear2():
"""linear2 -> single_transformer_blocks.X.proj_out (no chunking)."""
net = _load_via(C.try_load_lora, sd_lora_bfl_single_linear2())
assert net is not None and len(net.modules) == 1
assert 'lora_transformer_single_transformer_blocks_0_proj_out' in net.modules
return True
def test_lora_kohya_img_attn_proj():
"""Kohya flat-underscore on non-fused target binds with same diffusers-path key as BFL."""
net = _load_via(C.try_load_lora, sd_lora_kohya_img_attn_proj())
assert net is not None and len(net.modules) == 1
assert 'lora_transformer_transformer_blocks_0_attn_to_out_0' in net.modules
return True
def test_lora_kohya_img_attn_qkv_chunked():
"""Kohya fused img_attn.qkv splits into 3 chunks same as BFL form."""
net = _load_via(C.try_load_lora, sd_lora_kohya_img_attn_qkv_fused())
assert net is not None and len(net.modules) == 3
expected = {
'lora_transformer_transformer_blocks_0_attn_to_q',
'lora_transformer_transformer_blocks_0_attn_to_k',
'lora_transformer_transformer_blocks_0_attn_to_v',
}
assert set(net.modules) == expected
return True
def test_lora_peft_to_q():
"""PEFT format with diffusers paths passes through unchanged."""
net = _load_via(C.try_load_lora, sd_lora_peft_to_q())
assert net is not None and len(net.modules) == 1
assert 'lora_transformer_transformer_blocks_0_attn_to_q' in net.modules
return True
def test_lora_onetrainer_diffusers_flat():
"""OneTrainer lora_transformer_ diffusers-flat keys load via passthrough."""
net = _load_via(C.try_load_lora, sd_lora_onetrainer_diffusers_flat())
assert net is not None and len(net.modules) == 5, f'got {net.modules if net else None}'
expected = {
'lora_transformer_single_transformer_blocks_0_attn_to_q',
'lora_transformer_transformer_blocks_0_attn_to_out_0',
'lora_transformer_transformer_blocks_0_attn_add_k_proj',
'lora_transformer_transformer_blocks_0_ff_net_0_proj',
'lora_transformer_transformer_blocks_0_ff_net_2',
}
assert set(net.modules) == expected, f'got {set(net.modules)}'
return True
def test_lora_distilled_guidance():
"""LoRA on distilled_guidance_layer passes through unchanged (real module path)."""
net = _load_via(C.try_load_lora, sd_lora_distilled_guidance())
assert net is not None and len(net.modules) == 1
assert 'lora_transformer_distilled_guidance_layer_in_proj' in net.modules
return True
def sd_lokr_bfl_extra_modules():
"""BFL LoKR spanning the non-block extra targets and guidance MLP leaves.
Full-matrix factors with the ai-toolkit placeholder alpha, mirroring the
layout of real full-preset checkpoints.
"""
bases = [
'img_in', 'txt_in', 'final_layer.linear',
'distilled_guidance_layer.layers.0.in_layer',
'distilled_guidance_layer.layers.1.out_layer',
]
sd = {}
for base in bases:
sd[f'diffusion_model.{base}.lokr_w1'] = torch.randn(LOKR_W1_DIM, LOKR_W1_DIM)
sd[f'diffusion_model.{base}.lokr_w2'] = torch.randn(HIDDEN // LOKR_W1_DIM, HIDDEN // LOKR_W1_DIM)
sd[f'diffusion_model.{base}.alpha'] = torch.tensor(9999220736.0)
return sd
def test_lokr_bfl_extra_and_guidance():
"""Embedder/final-layer renames and guidance MLP leaf renames all bind."""
net = _load_via(C.try_load_lokr, sd_lokr_bfl_extra_modules())
assert net is not None and len(net.modules) == 5, f'got {sorted(net.modules) if net else None}'
expected = {
'lora_transformer_x_embedder',
'lora_transformer_context_embedder',
'lora_transformer_proj_out',
'lora_transformer_distilled_guidance_layer_layers_0_linear_1',
'lora_transformer_distilled_guidance_layer_layers_1_linear_2',
}
assert set(net.modules) == expected, f'got {set(net.modules)}'
for nk, mod in net.modules.items():
assert mod.dim is None and mod.calc_scale() == 1.0, f'{nk}: dim={mod.dim} scale={mod.calc_scale()}'
return True
def test_full_diff_chain():
"""Full-diff extraction loads through the chain; fused qkv diff skips."""
sd = {
'diffusion_model.double_blocks.0.img_attn.proj.diff': torch.randn(HIDDEN, HIDDEN),
'diffusion_model.double_blocks.0.img_attn.proj.diff_b': torch.randn(HIDDEN),
'diffusion_model.double_blocks.0.img_attn.qkv.diff': torch.randn(3 * HIDDEN, HIDDEN),
}
net = _load_via(C.try_load, sd)
assert net is not None and len(net.modules) == 1, f'got {net.modules if net else None}'
assert 'lora_transformer_transformer_blocks_0_attn_to_out_0' in net.modules, f'got {set(net.modules)}'
mod = next(iter(net.modules.values()))
updown, ex_bias = mod.calc_updown(mod.sd_module.weight)
assert tuple(updown.shape) == (HIDDEN, HIDDEN) and torch.isfinite(updown).all()
assert ex_bias is not None and tuple(ex_bias.shape) == (HIDDEN,)
return True
def test_lora_dora_threading():
"""dora_scale flows into NetworkModuleLora."""
net = _load_via(C.try_load_lora, sd_lora_with_dora_scale())
assert net is not None and len(net.modules) == 1
mod = next(iter(net.modules.values()))
assert mod.dora_scale is not None
return True
def test_lokr_bfl_img_attn_proj():
"""BFL LoKR on non-fused proj binds via NetworkModuleLokr (no chunk class)."""
net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_proj())
assert net is not None and len(net.modules) == 1
assert 'lora_transformer_transformer_blocks_0_attn_to_out_0' in net.modules
mod = next(iter(net.modules.values()))
assert isinstance(mod, network_lokr.NetworkModuleLokr) and not isinstance(mod, network_lokr.NetworkModuleLokrChunk)
return True
def test_lokr_bfl_img_attn_qkv_chunked():
"""BFL LoKR on fused img_attn.qkv emits 3 LokrChunk modules (equal chunks)."""
net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_qkv_equal_chunks())
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
expected = {
'lora_transformer_transformer_blocks_0_attn_to_q',
'lora_transformer_transformer_blocks_0_attn_to_k',
'lora_transformer_transformer_blocks_0_attn_to_v',
}
assert set(net.modules) == expected
for nk, mod in net.modules.items():
assert isinstance(mod, network_lokr.NetworkModuleLokrChunk), f'{nk}: type={type(mod).__name__}'
assert mod.num_chunks == 3, f'{nk}: num_chunks={mod.num_chunks}'
return True
def test_lokr_bfl_single_linear1_unequal_chunks():
"""BFL LoKR on fused linear1 emits 4 SliceChunks with UNEQUAL ranges.
Critical chroma-specific path: HIDDEN/HIDDEN/HIDDEN/MLP_HIDDEN partition.
"""
net = _load_via(C.try_load_lokr, sd_lokr_bfl_single_linear1_unequal())
assert net is not None and len(net.modules) == 4, f'got {net.modules if net else None}'
proj_mlp_key = 'lora_transformer_single_transformer_blocks_0_proj_mlp'
assert proj_mlp_key in net.modules
proj_mlp = net.modules[proj_mlp_key]
assert proj_mlp.end_row - proj_mlp.start_row == MLP_HIDDEN, \
f'proj_mlp range={proj_mlp.start_row}:{proj_mlp.end_row}, expected width={MLP_HIDDEN}'
for proj in ('attn_to_q', 'attn_to_k', 'attn_to_v'):
nk = f'lora_transformer_single_transformer_blocks_0_{proj}'
mod = net.modules[nk]
assert mod.end_row - mod.start_row == HIDDEN, f'{nk}: range={mod.start_row}:{mod.end_row}'
return True
def test_loha_bfl_img_attn_proj():
"""LoHA on non-fused target binds via NetworkModuleHada."""
net = _load_via(C.try_load_loha, sd_loha_bfl_img_attn_proj())
assert net is not None and len(net.modules) == 1
mod = next(iter(net.modules.values()))
assert isinstance(mod, network_hada.NetworkModuleHada)
return True
def test_loha_bfl_img_attn_qkv_chunked():
"""LoHA on fused img_attn.qkv emits 3 HadaChunk modules (equal chunks)."""
net = _load_via(C.try_load_loha, sd_loha_bfl_img_attn_qkv_skipped())
assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}'
expected = {
'lora_transformer_transformer_blocks_0_attn_to_q',
'lora_transformer_transformer_blocks_0_attn_to_k',
'lora_transformer_transformer_blocks_0_attn_to_v',
}
assert set(net.modules) == expected
for mod in net.modules.values():
assert isinstance(mod, network_hada.NetworkModuleHadaChunk)
return True
def test_oft_bfl_img_attn_proj():
"""LyCORIS oft_diag form loads on non-fused target without NoneType errors."""
net = _load_via(C.try_load_oft, sd_oft_bfl_img_attn_proj())
assert net is not None and len(net.modules) == 1
mod = next(iter(net.modules.values()))
assert isinstance(mod, network_oft.NetworkModuleOFT)
return True
def test_oft_bfl_img_attn_qkv_skipped():
"""OFT on fused img_attn.qkv is dropped (no row-sliceable OFT structure)."""
net = _load_via(C.try_load_oft, sd_oft_bfl_img_attn_qkv_skipped())
assert net is None or len(net.modules) == 0
return True
CAT_MATH = category('math')
def test_lora_calc_updown_shape():
net = _load_via(C.try_load_lora, sd_lora_bfl_img_attn_proj())
mod = make_network_for_module(next(iter(net.modules.values())))
target = torch.randn(HIDDEN, HIDDEN)
updown, _ = mod.calc_updown(target)
assert_shape(updown, target.shape, label='LoRA calc_updown')
return True
def test_lokr_calc_updown_shape():
net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_proj())
mod = make_network_for_module(next(iter(net.modules.values())))
target = torch.randn(HIDDEN, HIDDEN)
updown, _ = mod.calc_updown(target)
assert_shape(updown, target.shape, label='LoKR calc_updown')
return True
def test_lokr_chunk_equal_calc_updown_shape():
"""LokrChunk equal-chunks dispatch produces (HIDDEN, HIDDEN) output for the QKV split."""
net = _load_via(C.try_load_lokr, sd_lokr_bfl_img_attn_qkv_equal_chunks())
mod = make_network_for_module(next(iter(net.modules.values())))
target = torch.randn(HIDDEN, HIDDEN)
updown, _ = mod.calc_updown(target)
assert_shape(updown, target.shape, label='LokrChunk equal range')
return True
def test_lokr_slicechunk_unequal_calc_updown_shape():
"""LokrSliceChunk on the proj_mlp chunk produces (MLP_HIDDEN, HIDDEN) output.
This exercises the path that motivated NetworkModuleLokrSliceChunk's
existence: unequal partition where torch.chunk would not work.
"""
net = _load_via(C.try_load_lokr, sd_lokr_bfl_single_linear1_unequal())
proj_mlp_key = 'lora_transformer_single_transformer_blocks_0_proj_mlp'
mod = make_network_for_module(net.modules[proj_mlp_key])
target = torch.randn(MLP_HIDDEN, HIDDEN)
updown, _ = mod.calc_updown(target)
assert_shape(updown, target.shape, label='LokrSliceChunk unequal proj_mlp')
return True
def test_loha_calc_updown_shape():
net = _load_via(C.try_load_loha, sd_loha_bfl_img_attn_proj())
mod = make_network_for_module(next(iter(net.modules.values())))
target = torch.randn(HIDDEN, HIDDEN)
updown, _ = mod.calc_updown(target)
assert_shape(updown, target.shape, label='LoHA calc_updown')
return True
def test_oft_calc_updown_shape():
net = _load_via(C.try_load_oft, sd_oft_bfl_img_attn_proj())
mod = make_network_for_module(next(iter(net.modules.values())))
target = torch.randn(HIDDEN, HIDDEN)
updown, _ = mod.calc_updown(target)
assert_shape(updown, target.shape, label='OFT calc_updown')
return True
def run_tests():
t0 = time.time()
log.warning('=== Parsing primitives ===')
for fn in [test_parse_key_all_prefixes, test_marker_disambiguation, test_resolve_targets_static_renames,
test_resolve_targets_extra_and_guidance,
test_resolve_targets_onetrainer_passthrough]:
run_test(CAT_PARSE, fn)
log.warning('=== Loaders ===')
for fn in [
test_lora_bfl_img_attn_proj,
test_lora_bfl_img_attn_qkv_chunked,
test_lora_bfl_txt_attn_qkv_chunked,
test_lora_bfl_img_mlp,
test_lora_bfl_txt_mlp,
test_lora_bfl_single_linear1_unequal_chunks,
test_lora_bfl_single_linear2,
test_lora_kohya_img_attn_proj,
test_lora_kohya_img_attn_qkv_chunked,
test_lora_peft_to_q,
test_lora_onetrainer_diffusers_flat,
test_lora_distilled_guidance,
test_lora_dora_threading,
test_lokr_bfl_img_attn_proj,
test_lokr_bfl_img_attn_qkv_chunked,
test_lokr_bfl_single_linear1_unequal_chunks,
test_lokr_bfl_extra_and_guidance,
test_full_diff_chain,
test_loha_bfl_img_attn_proj,
test_loha_bfl_img_attn_qkv_chunked,
test_oft_bfl_img_attn_proj,
test_oft_bfl_img_attn_qkv_skipped,
]:
run_test(CAT_LOADER, fn)
log.warning('=== calc_updown shape sanity ===')
for fn in [
test_lora_calc_updown_shape,
test_lokr_calc_updown_shape,
test_lokr_chunk_equal_calc_updown_shape,
test_lokr_slicechunk_unequal_calc_updown_shape,
test_loha_calc_updown_shape,
test_oft_calc_updown_shape,
]:
run_test(CAT_MATH, fn)
elapsed = time.time() - t0
log.warning('=== Results ===')
total_pass = 0
total_fail = 0
for cat, info in results.items():
status = 'PASS' if info['failed'] == 0 else 'FAIL'
log.info(f' {cat}: {info["passed"]} passed, {info["failed"]} failed [{status}]')
total_pass += info['passed']
total_fail += info['failed']
log.warning(f'Total: {total_pass} passed, {total_fail} failed in {elapsed:.2f}s')
return total_fail == 0
if __name__ == '__main__':
ok = run_tests()
sys.exit(0 if ok else 1)