add Flux2/Klein LoRA support

- detect f2 model type for LoRAs via metadata, architecture, and filename/folder
- preprocess bare BFL-format keys with diffusion_model prefix for Flux2LoraLoaderMixin
- handle LoKR format via native NetworkModuleLokr with on-the-fly kron(w1, w2)
- add NetworkModuleLokrChunk for fused QKV split into separate Q/K/V modules
- activate native modules loaded via diffusers path
- improve error message for Flux1/Flux2 architecture mismatch

modules/lora/lora_load.py

@@ -257,7 +257,11 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
                     shared.compiled_model_state.lora_model.append(f"{name}:{lora_scale}")
                 lora_method = lora_overrides.get_method(shorthash)
                 if lora_method == 'diffusers':
-                    net = lora_diffusers.load_diffusers(name, network_on_disk, lora_scale, lora_module)
+                    if shared.sd_model_type == 'f2':
+                        from pipelines.flux import flux2_lora
+                        net = flux2_lora.try_load_lokr(name, network_on_disk, lora_scale)
+                    if net is None:
+                        net = lora_diffusers.load_diffusers(name, network_on_disk, lora_scale, lora_module)
                 elif lora_method == 'nunchaku':
                     pass # handled directly from extra_networks_lora.load_nunchaku
                 else:
@@ -272,7 +276,11 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
                 continue
         if net is None:
             failed_to_load_networks.append(name)
-            log.error(f'Network load: type=LoRA name="{name}" detected={network_on_disk.sd_version if network_on_disk is not None else None} not found')
+            lora_ver = network_on_disk.sd_version if network_on_disk is not None else None
+            if lora_ver in ('f1', '') and shared.sd_model_type == 'f2':
+                log.error(f'Network load: type=LoRA name="{name}" incompatible: Flux1 LoRA cannot be used with Flux2/Klein')
+            else:
+                log.error(f'Network load: type=LoRA name="{name}" detected={lora_ver} not found')
             continue
         if hasattr(sd_model, 'embedding_db'):
             sd_model.embedding_db.load_diffusers_embedding(None, net.bundle_embeddings)
@@ -309,6 +317,12 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
                 errors.display(e, 'LoRA')
         shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, force=True, silent=True) # some layers may end up on cpu without hook
 
+    # Activate native modules loaded via diffusers path (e.g., LoKR on Flux2)
+    native_nets = [net for net in l.loaded_networks if len(net.modules) > 0]
+    if native_nets:
+        from modules.lora import networks
+        networks.network_activate()
+
     if len(l.loaded_networks) > 0 and l.debug:
         log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}')
 

modules/lora/network.py

@@ -58,6 +58,8 @@ class NetworkOnDisk:
             return 'sc'
         if base.startswith("sd3"):
             return 'sd3'
+        if base.startswith("flux2") or "klein" in base:
+            return 'f2'
         if base.startswith("flux"):
             return 'f1'
         if base.startswith("hunyuan_video"):
@@ -75,6 +77,8 @@ class NetworkOnDisk:
             return 'xl'
         if arch.startswith("stable-cascade"):
             return 'sc'
+        if arch.startswith("flux2") or "klein" in arch:
+            return 'f2'
         if arch.startswith("flux"):
             return 'f1'
         if arch.startswith("hunyuan-video"):
@@ -86,6 +90,8 @@ class NetworkOnDisk:
             return 'sd1'
         if str(self.metadata.get('ss_v2', "")) == "True":
             return 'sd2'
+        if 'klein' in self.name.lower() or 'klein' in self.fullname.lower():
+            return 'f2'
         if 'flux' in self.name.lower():
             return 'f1'
         if 'xl' in self.name.lower():

modules/lora/network_lokr.py

@@ -55,3 +55,40 @@ class NetworkModuleLokr(network.NetworkModule): # pylint: disable=abstract-metho
             output_shape = target.shape
         updown = make_kron(output_shape, w1, w2)
         return self.finalize_updown(updown, target, output_shape)
+
+
+class NetworkModuleLokrChunk(NetworkModuleLokr):
+    """LoKR module that returns one chunk of the Kronecker product.
+
+    Used when a LoKR adapter targets a fused weight (e.g., QKV) but the model
+    has separate modules (Q, K, V). Computes kron(w1, w2) on-the-fly and
+    returns only the designated chunk, keeping memory usage minimal.
+    """
+    def __init__(self, net, weights, chunk_index, num_chunks):
+        super().__init__(net, weights)
+        self.chunk_index = chunk_index
+        self.num_chunks = num_chunks
+
+    def calc_updown(self, target):
+        if self.w1 is not None:
+            w1 = self.w1.to(target.device, dtype=target.dtype)
+        else:
+            w1a = self.w1a.to(target.device, dtype=target.dtype)
+            w1b = self.w1b.to(target.device, dtype=target.dtype)
+            w1 = w1a @ w1b
+        if self.w2 is not None:
+            w2 = self.w2.to(target.device, dtype=target.dtype)
+        elif self.t2 is None:
+            w2a = self.w2a.to(target.device, dtype=target.dtype)
+            w2b = self.w2b.to(target.device, dtype=target.dtype)
+            w2 = w2a @ w2b
+        else:
+            t2 = self.t2.to(target.device, dtype=target.dtype)
+            w2a = self.w2a.to(target.device, dtype=target.dtype)
+            w2b = self.w2b.to(target.device, dtype=target.dtype)
+            w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
+        full_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)]
+        updown = make_kron(full_shape, w1, w2)
+        updown = torch.chunk(updown, self.num_chunks, dim=0)[self.chunk_index]
+        output_shape = list(updown.shape)
+        return self.finalize_updown(updown, target, output_shape)