From 72511f1bd76480b16db4edc95bc01b987bc3bc59 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 13 Jul 2026 00:00:10 +0100 Subject: [PATCH] feat(lora): reconstruct the lycoris sparse bias residual triplet LyCORIS extraction with use_sparse_bias saves bias_indices/bias_values/ bias_size per module: the sparse weight-shaped remainder of the SVD extraction, named bias for historical reasons. The keys were dropped by both loader paths, so extracted adapters applied without the residual correction; the dense-bias branch in finalize_updown that consumes it was unreachable. - rebuild the COO tensor in NetworkModule.__init__ (int16 indices cast to long), shared by the native and generic loaders; kept sparse so the dense += sparse in finalize_updown materializes per module at apply instead of near-model-size densification at load - accept the triplet suffixes in LORA_SUFFIXES; fused targets skip with the weight-shaped-bias warning - cover an extraction-faithful numeric round-trip and the fused skip in the offline suite --- modules/lora/native_adapter.py | 10 +++--- modules/lora/network.py | 11 +++++++ test/test-flux2-native-adapters.py | 50 ++++++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 4 deletions(-) diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index 9b34429c0..d6dc8537e 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -104,8 +104,10 @@ LORA_SUFFIXES = ( # diff_b: bias delta some saves pair with the weight LoRA, applied as ex_bias. # magnitude / lora_magnitude_vector: DoRA row norms (ai-toolkit / PEFT key # names); converted onto the dora_scale path by try_load_lora. + # bias_indices/values/size: LyCORIS extraction sparse residual triplet, + # reconstructed into the dense-bias path by network.NetworkModule. ".alpha", ".dora_scale", ".magnitude", ".lora_magnitude_vector", - ".bias", ".diff_b", ".scale", + ".bias", ".bias_indices", ".bias_values", ".bias_size", ".diff_b", ".scale", ) LOKR_SUFFIXES = ( ".lokr_w1", ".lokr_w2", @@ -586,9 +588,9 @@ def try_load_lora(name, network_on_disk, lora_scale, *, target_w = w if chunk is not None: - if "bias" in w: - # Legacy weight-shaped bias (LyCORIS sparse-residual heritage) - # has no defined partition on a fused target. + if "bias" in w or "bias_indices" in w: + # Weight-shaped bias residuals (dense or LyCORIS sparse + # triplet) are not partitioned onto fused targets. log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} weight-shaped bias on fused target skipped (unsupported)') skipped += 1 continue diff --git a/modules/lora/network.py b/modules/lora/network.py index 38a0ebc3c..e05f6abe9 100644 --- a/modules/lora/network.py +++ b/modules/lora/network.py @@ -1,6 +1,7 @@ import os import enum from collections import namedtuple +import torch from modules import hashes, shared, sd_checkpoint @@ -174,6 +175,16 @@ class NetworkModule: self.shape = self.sd_module.weight.shape self.dim = None self.bias = weights.w.get("bias") + if self.bias is None and "bias_indices" in weights.w: + # LyCORIS extraction with use_bias: the sparse weight-shaped + # remainder of the SVD extraction ("bias" is historical naming), + # stored COO with int16 indices. Kept sparse; finalize_updown's + # dense += sparse materializes it per module at apply time. + self.bias = torch.sparse_coo_tensor( + weights.w["bias_indices"].to(torch.long), + weights.w["bias_values"], + tuple(weights.w["bias_size"]), + ) self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None self.scale = weights.w["scale"].item() if "scale" in weights.w else None self.dora_scale = weights.w.get("dora_scale", None) diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py index cf4199add..1375e25a1 100644 --- a/test/test-flux2-native-adapters.py +++ b/test/test-flux2-native-adapters.py @@ -1071,6 +1071,54 @@ def test_lora_peft_magnitude_vector(): 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() @@ -1560,6 +1608,8 @@ def run_tests(): 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,