From 726d3cf06faa3e0bd8b686d09fca9598c69ecbb1 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Tue, 23 Jun 2026 22:02:34 +0200 Subject: [PATCH] fix(arch): detect all Krea 2 LoRA forms via the attention gate marker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- lora_optimizer.py | 27 ++++++++++++++------------- pyproject.toml | 2 +- tests/test_lora_optimizer.py | 21 +++++++++++++++++++++ 3 files changed, 36 insertions(+), 14 deletions(-) diff --git a/lora_optimizer.py b/lora_optimizer.py index 4d0f94a..04c0857 100644 --- a/lora_optimizer.py +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 78ef09b..0944018 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/test_lora_optimizer.py b/tests/test_lora_optimizer.py index 3d8dda4..aa81b2b 100644 --- a/tests/test_lora_optimizer.py +++ b/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",