fix(arch): detect all Krea 2 LoRA forms via the attention gate marker

The initial krea2 detection only matched the Comfy-Org native form
(blocks.N.attn.wq). Real trainer outputs use diffusers-style keys and were
misdetected: the 'krea_2' trainer (diffusion_model.transformer_blocks.N.attn
.to_q/to_gate) hit the ACE-Step branch, and the diffusers form
(transformer.transformer_blocks.N...) would hit qwen.

The marker common to ALL forms is the per-attention sigmoid GATE
(attn.to_gate / attn.gate / attn.w{q,k,v,o}) — unique among supported archs.
Detect on that, moved BEFORE the qwen/flux/ACE-Step branches. Verified vs 3
real LoRAs (official Comfy-Org rank-64, a krea_2-trained style LoRA, and a
diffusers 'softwatercolor' with text_fusion) — all -> krea2/dit. +2 tests
(ACE-Step/qwen collision guards); 284 pass. Registry 1.7.12.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

lora_optimizer.py

@@ -1046,6 +1046,20 @@ class _LoRAMergeBase:
                 if shape is not None and len(shape) >= 1 and shape[0] == 13824:
                     return 'ideogram4'
 
+        # Krea 2 (krea/Krea-2) — from-scratch single-stream image DiT (NOT the
+        # older FLUX.1-Krea finetune). The unique marker across every trainer
+        # form is a per-attention sigmoid GATE projection — attn.to_gate (the
+        # "krea_2" trainer + diffusers forms) / attn.gate / attn.w{q,k,v,o} (the
+        # Comfy-Org native form). No other supported arch has an attention gate,
+        # so this MUST run before the qwen (transformer.transformer_blocks), flux,
+        # and ACE-Step (transformer_blocks.N.attn.to_q) branches would otherwise
+        # claim its diffusers-style keys. Verified vs 3 real LoRA forms:
+        #   diffusion_model.blocks.N.attn.wq            (Comfy-Org native)
+        #   diffusion_model.transformer_blocks.N.attn.to_gate  (krea_2 trainer)
+        #   transformer.transformer_blocks.N.attn.to_gate + text_fusion (diffusers)
+        if any(re.search(r'attn[._](?:to_gate|gate|w[qkvo])(?=[._])', k) for k in keys):
+            return 'krea2'
+
         # Z-Image Turbo (Lumina2): layers.N with attention patterns
         # Handles: diffusion_model.layers.N, single_transformer_blocks.N (non-FLUX),
         #          lora_unet_layers_N (Musubi Tuner)
@@ -1090,19 +1104,6 @@ class _LoRAMergeBase:
         if any('double_blocks' in k or 'single_blocks' in k for k in keys):
             return 'flux'
 
-        # Krea 2 (krea/Krea-2) — a from-scratch 12B single-stream image DiT, NOT
-        # the older FLUX.1-Krea finetune. diffusion_model.blocks.N with GQA
-        # attention named attn.wq/wk/wv/wo (wk/wv narrower than wq under GQA), a
-        # sigmoid attn.gate, and SwiGLU mlp.gate/up/down. The attn.w{q,k,v,o}
-        # naming is unique — WAN uses self_attn.q/k/v/o + ffn — so this is safe
-        # before the WAN/Anima single-stream blocks.N checks. Verified against the
-        # official Comfy-Org/Krea-2 rank-64 LoRA.
-        if any(re.search(r'blocks[._]\d+[._]attn[._]w[qkvo](?=[._])', k) for k in keys):
-            return 'krea2'
-        if (any(re.search(r'blocks[._]\d+[._]attn[._]gate(?=[._])', k) for k in keys)
-                and any(re.search(r'blocks[._]\d+[._]mlp[._]gate(?=[._])', k) for k in keys)):
-            return 'krea2'
-
         # Anima (CircleStone Labs) — a Cosmos-Predict2 DiT with SPLIT QKV. MUST be
         # checked before ACE-Step / Wan / LTX: its blocks.N.self_attn.q_proj and
         # its Qwen3 text encoder (lora_te_layers_N_*) otherwise match those.

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-lora-optimizer"
 description = "Auto-optimizer node for combining multiple LoRAs via diff-based merging with TIES conflict resolution"
-version = "1.7.11"
+version = "1.7.12"
 license = { text = "GPL-3.0-only" }
 dependencies = [
     "scikit-learn>=1.3",

tests/test_lora_optimizer.py

@@ -4624,6 +4624,27 @@ class Krea2DetectionTests(unittest.TestCase):
                  "lora_unet_blocks_0_mlp_gate.lora_up.weight"])
         self.assertEqual(self.det(s), "krea2")
 
+    def test_trainer_diffusion_model_form(self):
+        # The "krea_2" trainer: diffusion_model.transformer_blocks.N.attn.to_*
+        # with a sigmoid attn.to_gate. Must NOT be mistaken for ACE-Step (which
+        # also matches transformer_blocks.N.attn.to_q).
+        s = _sd(["diffusion_model.transformer_blocks.0.attn.to_q.lora_A.weight",
+                 "diffusion_model.transformer_blocks.0.attn.to_k.lora_A.weight",
+                 "diffusion_model.transformer_blocks.0.attn.to_v.lora_A.weight",
+                 "diffusion_model.transformer_blocks.0.attn.to_out.0.lora_A.weight",
+                 "diffusion_model.transformer_blocks.0.attn.to_gate.lora_A.weight",
+                 "diffusion_model.transformer_blocks.0.ff.gate.lora_A.weight"])
+        self.assertEqual(self.det(s), "krea2")
+
+    def test_diffusers_transformer_form(self):
+        # diffusers form: transformer.transformer_blocks.N.attn.to_* + to_gate
+        # + text_fusion. Must NOT be mistaken for Qwen-Image (transformer.transformer_blocks).
+        s = _sd(["transformer.transformer_blocks.0.attn.to_q.lora_A.weight",
+                 "transformer.transformer_blocks.0.attn.to_gate.lora_A.weight",
+                 "transformer.transformer_blocks.0.ff.gate.lora_A.weight",
+                 "transformer.text_fusion.refiner_blocks.0.attn.to_gate.lora_A.weight"])
+        self.assertEqual(self.det(s), "krea2")
+
     def test_gate_plus_mlp_gate_fallback(self):
         # Backup discriminator when attn.w{q,k,v,o} isn't in a partial LoRA.
         s = _sd(["diffusion_model.blocks.3.attn.gate.lora_up.weight",