Add three-tier LoKR quality comparison (fuse-first, Kronecker split, SVD)

- Add fuse_qkv parameter to BFL LoKR converter for lossless fuse-first path
- Thread fuse_qkv through lora_pipeline.py (lora_state_dict -> load_lora_weights)
- Fuse model QKV projections before adapter injection when fuse_qkv=True
- Update benchmark script with --tiers, --no-offload flags for all three paths

Author: CalamitousFelicitousness

benchmark_lokr.py

@@ -1,12 +1,20 @@
-"""Benchmark: Lossless LoKR vs Lossy LoRA-via-SVD on Flux2 Klein 9B.
+"""Benchmark: Three-tier LoKR quality comparison on Flux2 Klein 9B.
+
+Tier 1 - Fuse-first (lossless): Fuse model QKV, map BFL LoKR directly. Exact.
+Tier 2 - Kronecker split (default): Split fused QKV via Van Loan re-factorization. Slight loss.
+Tier 3 - SVD to LoRA (fully lossy): Convert entire LoKR to LoRA via peft.convert_to_lora.
+
+Tiers 1+2 only apply to BFL-format LoKR (fused QKV). LyCORIS and diffusers-native
+formats already have separate Q/K/V and only run the default path.
 
-Generates images using both conversion paths for visual comparison.
 Uses bf16 with CPU offload.
 
 Usage:
     python benchmark_lokr.py
     python benchmark_lokr.py --lokr-path "puttmorbidly233/lora" --lokr-name "klein_snofs_v1_2.safetensors"
     python benchmark_lokr.py --prompt "a portrait in besch art style" --ranks 32 64 128
+    python benchmark_lokr.py --tiers 1 2     # skip SVD tier
+    python benchmark_lokr.py --tiers 2 3     # skip fuse-first tier
 """
 
 import argparse
@@ -15,18 +23,22 @@
 import time
 
 import torch
+
 from diffusers import Flux2KleinPipeline
-from peft import convert_to_lora
+
 
 MODEL_ID = "black-forest-labs/FLUX.2-klein-9B"
 DEFAULT_LOKR_PATH = "gattaplayer/besch-flux2-klein-9b-lokr-lion-3e-6-bs2-ga2-v02"
 OUTPUT_DIR = "benchmark_output"
 
 
-def load_pipeline():
-    """Load Flux2 Klein 9B in bf16 with model CPU offload."""
+def load_pipeline(no_offload=False):
+    """Load Flux2 Klein 9B in bf16."""
     pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
-    pipe.enable_model_cpu_offload()
+    if no_offload:
+        pipe = pipe.to("cuda")
+    else:
+        pipe.enable_model_cpu_offload()
     return pipe
 
 
@@ -44,9 +56,34 @@ def generate(pipe, prompt, seed, num_steps=4, guidance_scale=1.0):
     return image
 
 
-def benchmark_lossless(pipe, prompt, seed, lokr_path, lokr_name):
-    """Path A: Load LoKR natively (lossless)."""
-    print("\n=== Path A: Lossless LoKR ===")
+def benchmark_baseline(pipe, prompt, seed):
+    """Baseline: No adapter."""
+    print("\n=== Baseline: No adapter ===")
+    t0 = time.time()
+    image = generate(pipe, prompt, seed)
+    print(f"  Generated in {time.time() - t0:.1f}s")
+    return image
+
+
+def benchmark_tier1_fuse_first(pipe, prompt, seed, lokr_path, lokr_name):
+    """Tier 1: Fuse model QKV, then load BFL LoKR directly (lossless)."""
+    print("\n=== Tier 1: Fuse-first LoKR (lossless) ===")
+    t0 = time.time()
+    kwargs = {"weight_name": lokr_name} if lokr_name else {}
+    pipe.load_lora_weights(lokr_path, fuse_qkv=True, **kwargs)
+    print(f"  Loaded in {time.time() - t0:.1f}s")
+
+    t0 = time.time()
+    image = generate(pipe, prompt, seed)
+    print(f"  Generated in {time.time() - t0:.1f}s")
+
+    pipe.unload_lora_weights()
+    return image
+
+
+def benchmark_tier2_kronecker_split(pipe, prompt, seed, lokr_path, lokr_name):
+    """Tier 2: Split fused QKV via Kronecker re-factorization (default path)."""
+    print("\n=== Tier 2: Kronecker split LoKR (default) ===")
     t0 = time.time()
     kwargs = {"weight_name": lokr_name} if lokr_name else {}
     pipe.load_lora_weights(lokr_path, **kwargs)
@@ -60,15 +97,16 @@ def benchmark_lossless(pipe, prompt, seed, lokr_path, lokr_name):
     return image
 
 
-def benchmark_lossy(pipe, prompt, seed, rank, lokr_path, lokr_name):
-    """Path B: Load LoKR, convert to LoRA via SVD (lossy)."""
-    print(f"\n=== Path B: Lossy LoRA via SVD (rank={rank}) ===")
+def benchmark_tier3_svd(pipe, prompt, seed, rank, lokr_path, lokr_name):
+    """Tier 3: Convert LoKR to LoRA via SVD (fully lossy)."""
+    from peft import convert_to_lora, inject_adapter_in_model, set_peft_model_state_dict
+
+    print(f"\n=== Tier 3: SVD to LoRA (rank={rank}) ===")
     t0 = time.time()
     kwargs = {"weight_name": lokr_name} if lokr_name else {}
     pipe.load_lora_weights(lokr_path, **kwargs)
     load_time = time.time() - t0
 
-    # Detect the actual adapter name assigned by peft
     adapter_name = next(iter(pipe.transformer.peft_config.keys()))
     print(f"  Adapter name: {adapter_name}")
 
@@ -77,9 +115,6 @@ def benchmark_lossy(pipe, prompt, seed, rank, lokr_path, lokr_name):
     convert_time = time.time() - t0
     print(f"  Loaded LoKR in {load_time:.1f}s, converted to LoRA in {convert_time:.1f}s")
 
-    # Replace LoKR adapter with converted LoRA
-    from peft import inject_adapter_in_model, set_peft_model_state_dict
-
     pipe.transformer.delete_adapters(adapter_name)
     inject_adapter_in_model(lora_config, pipe.transformer, adapter_name=adapter_name)
     set_peft_model_state_dict(pipe.transformer, lora_sd, adapter_name=adapter_name)
@@ -92,24 +127,18 @@ def benchmark_lossy(pipe, prompt, seed, rank, lokr_path, lokr_name):
     return image
 
 
-def benchmark_baseline(pipe, prompt, seed):
-    """Baseline: No adapter."""
-    print("\n=== Baseline: No adapter ===")
-    t0 = time.time()
-    image = generate(pipe, prompt, seed)
-    print(f"  Generated in {time.time() - t0:.1f}s")
-    return image
-
-
 def main():
-    parser = argparse.ArgumentParser(description="Benchmark LoKR vs LoRA-via-SVD")
+    parser = argparse.ArgumentParser(description="Benchmark LoKR quality tiers")
     parser.add_argument("--prompt", default="a portrait painting in besch art style")
     parser.add_argument("--lokr-path", default=DEFAULT_LOKR_PATH, help="HF repo or local path to LoKR checkpoint")
     parser.add_argument("--lokr-name", default=None, help="Filename within HF repo (if multi-file)")
     parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--ranks", type=int, nargs="+", default=[32, 64, 128])
+    parser.add_argument(
+        "--tiers", type=int, nargs="+", default=[1, 2, 3], help="Tiers to run (1=fuse, 2=kronecker, 3=svd)"
+    )
+    parser.add_argument("--ranks", type=int, nargs="+", default=[32, 64, 128], help="SVD ranks for tier 3")
     parser.add_argument("--skip-baseline", action="store_true")
-    parser.add_argument("--skip-lossy", action="store_true")
+    parser.add_argument("--no-offload", action="store_true", help="Keep model on GPU instead of CPU offload")
     args = parser.parse_args()
 
     os.makedirs(OUTPUT_DIR, exist_ok=True)
@@ -118,11 +147,13 @@ def main():
     print(f"LoKR:  {args.lokr_path}" + (f" ({args.lokr_name})" if args.lokr_name else ""))
     print(f"Prompt: {args.prompt}")
     print(f"Seed: {args.seed}")
-    if not args.skip_lossy:
-        print(f"SVD ranks to test: {args.ranks}")
+    print(f"Tiers: {args.tiers}")
+    if 3 in args.tiers:
+        print(f"SVD ranks: {args.ranks}")
 
-    print("\nLoading pipeline (bf16, model CPU offload)...")
-    pipe = load_pipeline()
+    mode = "on GPU" if args.no_offload else "with CPU offload"
+    print(f"\nLoading pipeline (bf16, {mode})...")
+    pipe = load_pipeline(no_offload=args.no_offload)
 
     # Baseline
     if not args.skip_baseline:
@@ -131,28 +162,36 @@ def main():
         img.save(path)
         print(f"  Saved: {path}")
 
-    # Path A: Lossless LoKR
-    img = benchmark_lossless(pipe, args.prompt, args.seed, args.lokr_path, args.lokr_name)
-    path = os.path.join(OUTPUT_DIR, "lokr_lossless.png")
-    img.save(path)
-    print(f"  Saved: {path}")
+    # Tier 1: Fuse-first (lossless, BFL only)
+    if 1 in args.tiers:
+        img = benchmark_tier1_fuse_first(pipe, args.prompt, args.seed, args.lokr_path, args.lokr_name)
+        path = os.path.join(OUTPUT_DIR, "tier1_fuse_lossless.png")
+        img.save(path)
+        print(f"  Saved: {path}")
+        gc.collect()
+        torch.cuda.empty_cache()
 
-    gc.collect()
-    torch.cuda.empty_cache()
+    # Tier 2: Kronecker split (default)
+    if 2 in args.tiers:
+        img = benchmark_tier2_kronecker_split(pipe, args.prompt, args.seed, args.lokr_path, args.lokr_name)
+        path = os.path.join(OUTPUT_DIR, "tier2_kronecker.png")
+        img.save(path)
+        print(f"  Saved: {path}")
+        gc.collect()
+        torch.cuda.empty_cache()
 
-    # Path B: Lossy LoRA via SVD at various ranks
-    if not args.skip_lossy:
+    # Tier 3: SVD to LoRA at various ranks
+    if 3 in args.tiers:
         for rank in args.ranks:
-            img = benchmark_lossy(pipe, args.prompt, args.seed, rank, args.lokr_path, args.lokr_name)
-            path = os.path.join(OUTPUT_DIR, f"lora_svd_rank{rank}.png")
+            img = benchmark_tier3_svd(pipe, args.prompt, args.seed, rank, args.lokr_path, args.lokr_name)
+            path = os.path.join(OUTPUT_DIR, f"tier3_svd_rank{rank}.png")
             img.save(path)
             print(f"  Saved: {path}")
-
             gc.collect()
             torch.cuda.empty_cache()
 
     print(f"\nAll results saved to {OUTPUT_DIR}/")
-    print("Compare: baseline.png vs lokr_lossless.png vs lora_svd_rank*.png")
+    print("Compare: baseline.png vs tier1_fuse_lossless.png vs tier2_kronecker.png vs tier3_svd_rank*.png")
 
 
 if __name__ == "__main__":

src/diffusers/loaders/lora_conversion_utils.py

@@ -2688,11 +2688,15 @@ def _split_lokr_qkv(w1, w2, target_keys, factor):
     return result
 
 
-def _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict):
+def _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict, fuse_qkv=False):
     """Convert BFL-format Flux2 LoKR state dict to peft-compatible diffusers format.
 
-    Handles fused QKV by splitting via Kronecker re-factorization (Van Loan algorithm).
-    Non-QKV modules are remapped directly. Alpha scaling is baked into lokr_w1.
+    Args:
+        state_dict: BFL-format LoKR state dict with ``diffusion_model.`` prefix.
+        fuse_qkv: If True, map fused QKV directly to ``to_qkv``/``to_added_qkv`` targets
+            (lossless, but requires the model's QKV to be fused before injection).
+            If False (default), split fused QKV into separate Q/K/V via Kronecker
+            re-factorization (slightly lossy, no model fusion needed).
     """
     converted_state_dict = {}
 
@@ -2793,11 +2797,17 @@ def _remap_lokr_qkv(bfl_path, target_keys):
         tb = f"transformer_blocks.{dl}"
         db = f"double_blocks.{dl}"
 
-        # Split fused QKV into separate Q/K/V via Kronecker re-factorization
-        _remap_lokr_qkv(f"{db}.img_attn.qkv", [f"{tb}.attn.to_q", f"{tb}.attn.to_k", f"{tb}.attn.to_v"])
-        _remap_lokr_qkv(
-            f"{db}.txt_attn.qkv", [f"{tb}.attn.add_q_proj", f"{tb}.attn.add_k_proj", f"{tb}.attn.add_v_proj"]
-        )
+        if fuse_qkv:
+            # Lossless: map directly to fused targets (caller must fuse model QKV first)
+            _remap_lokr_module(f"{db}.img_attn.qkv", f"{tb}.attn.to_qkv")
+            _remap_lokr_module(f"{db}.txt_attn.qkv", f"{tb}.attn.to_added_qkv")
+        else:
+            # Split fused QKV into separate Q/K/V via Kronecker re-factorization
+            _remap_lokr_qkv(f"{db}.img_attn.qkv", [f"{tb}.attn.to_q", f"{tb}.attn.to_k", f"{tb}.attn.to_v"])
+            _remap_lokr_qkv(
+                f"{db}.txt_attn.qkv",
+                [f"{tb}.attn.add_q_proj", f"{tb}.attn.add_k_proj", f"{tb}.attn.add_v_proj"],
+            )
 
         # Projections
         _remap_lokr_module(f"{db}.img_attn.proj", f"{tb}.attn.to_out.0")

src/diffusers/loaders/lora_pipeline.py

@@ -5648,6 +5648,7 @@ def lora_state_dict(
         weight_name = kwargs.pop("weight_name", None)
         use_safetensors = kwargs.pop("use_safetensors", None)
         return_lora_metadata = kwargs.pop("return_lora_metadata", False)
+        fuse_qkv = kwargs.pop("fuse_qkv", False)
 
         allow_pickle = False
         if use_safetensors is None:
@@ -5691,14 +5692,16 @@ def lora_state_dict(
         is_lokr = any("lokr_" in k for k in state_dict)
         if is_lokr:
             if any(k.startswith("diffusion_model.") for k in state_dict):
-                state_dict = _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict)
+                state_dict = _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict, fuse_qkv=fuse_qkv)
             elif any(k.startswith("lycoris_") for k in state_dict):
                 state_dict = _convert_lycoris_flux2_lokr_to_diffusers(state_dict)
             else:
                 state_dict = _convert_diffusers_flux2_lokr_to_peft(state_dict)
             if metadata is None:
                 metadata = {}
             metadata["is_lokr"] = "true"
+            if fuse_qkv:
+                metadata["fuse_qkv"] = "true"
         else:
             is_ai_toolkit = any(k.startswith("diffusion_model.") for k in state_dict)
             if is_ai_toolkit:
@@ -5740,6 +5743,10 @@ def load_lora_weights(
 
         transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
 
+        # Fuse model QKV projections before injection if requested (lossless path for BFL LoKR)
+        if metadata and metadata.get("fuse_qkv") == "true":
+            transformer.fuse_qkv_projections()
+
         self.load_lora_into_transformer(
             state_dict,
             transformer=transformer,