From 26828757eb115112a3f3e201d63298762272a56e Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Tue, 23 Jun 2026 21:36:36 +0200 Subject: [PATCH] feat(arch): add Krea 2 detection + dit preset Krea 2 (krea/Krea-2) is a from-scratch 12B single-stream image DiT (NOT the older FLUX.1-Krea finetune) released 2026-06-22. Its LoRAs previously fell to 'unknown' -> sd_unet preset (generic, suboptimal merge tuning). Detect diffusion_model.blocks.N with GQA attn.wq/wk/wv/wo + sigmoid attn.gate + 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 the branch sits before WAN/Anima safely. Maps to the 'dit' preset like Flux/Z-Image; no key normalization needed (QKV already split, ComfyUI-native diffusion_model.* keys pass through). Verified against the official Comfy-Org/Krea-2 rank-64 LoRA (real keys -> arch=krea2 -> dit). +4 detection tests incl. WAN/flux/qwen collision guards; 282 pass. No algo bump (only affects krea2, which had no prior cache). Registry 1.7.11 (also publishes the accumulated tooltip fixes). Co-Authored-By: Claude Opus 4.8 --- lora_optimizer.py | 20 ++++++++++++++--- pyproject.toml | 2 +- tests/test_lora_optimizer.py | 42 ++++++++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 4 deletions(-) diff --git a/lora_optimizer.py b/lora_optimizer.py index 60b1a1a..4d0f94a 100644 --- a/lora_optimizer.py +++ b/lora_optimizer.py @@ -336,7 +336,7 @@ _ARCH_PRESETS = { "alignment_threshold": 0.1, "suggested_max_strength_cap": 5.0, "auto_strength_orthogonal_floor": 0.85, - "display_name": "DiT (Flux/WAN/Z-Image/LTX/Ideogram-4/HunyuanVideo)", + "display_name": "DiT (Flux/WAN/Z-Image/LTX/Ideogram-4/Krea-2/HunyuanVideo)", "full_rank": { "rank_threshold": 512, "disable_slerp_upgrade": True, @@ -405,7 +405,7 @@ _VIDEO_ARCH_ORTHOGONAL_FLOOR = {"wan": 1.0, "ltx": 1.0, "acestep": 1.0} _ARCH_TO_PRESET = { "sdxl": "sd_unet", "sd15": "sd_unet", "unknown": "sd_unet", "flux": "dit", "wan": "dit", "zimage": "dit", "ltx": "dit", - "ideogram4": "dit", "anima": "dit", + "ideogram4": "dit", "anima": "dit", "krea2": "dit", "acestep": "acestep_dit", "qwen_image": "llm", } @@ -1020,7 +1020,7 @@ class _LoRAMergeBase: """ Detect model architecture from LoRA key patterns. Returns: 'zimage', 'flux', 'wan', 'acestep', 'sdxl', 'sd15', 'ltx', - 'qwen_image', 'ideogram4', 'anima', or 'unknown'. + 'qwen_image', 'ideogram4', 'anima', 'krea2', or 'unknown'. """ keys = list(lora_sd.keys()) keys_str = ' '.join(k.lower() for k in keys) @@ -1090,6 +1090,19 @@ 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. @@ -6558,6 +6571,7 @@ class LoRAOptimizer(_LoRAMergeBase): 'ltx': 'LTX Video', 'qwen_image': 'Qwen-Image', 'anima': 'Anima (Cosmos-Predict2 DiT)', + 'krea2': 'Krea 2', } lines.append(f"Detected architecture: {arch_names.get(detected_arch, detected_arch)}") if normalize_keys == "enabled": diff --git a/pyproject.toml b/pyproject.toml index e00d868..78ef09b 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.10" +version = "1.7.11" 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 1bc11a0..3d8dda4 100644 --- a/tests/test_lora_optimizer.py +++ b/tests/test_lora_optimizer.py @@ -4601,6 +4601,48 @@ class AnimaDetectionTests(unittest.TestCase): self.assertEqual(self.det(ltx), "ltx") +class Krea2DetectionTests(unittest.TestCase): + """Krea 2 (krea/Krea-2, from-scratch single-stream image DiT) detection.""" + + det = staticmethod(lambda s: lora_optimizer.LoRAOptimizer._detect_architecture(s)) + + def test_comfy_native_form(self): + # Mirrors the official Comfy-Org/Krea-2 rank-64 LoRA: GQA attn.wq/wk/wv/wo + # + sigmoid attn.gate + SwiGLU mlp.gate/up/down under diffusion_model.blocks.N + s = _sd(["diffusion_model.blocks.0.attn.wq.lora_up.weight", + "diffusion_model.blocks.0.attn.wk.lora_up.weight", + "diffusion_model.blocks.0.attn.wv.lora_up.weight", + "diffusion_model.blocks.0.attn.wo.lora_up.weight", + "diffusion_model.blocks.0.attn.gate.lora_up.weight", + "diffusion_model.blocks.0.mlp.gate.lora_up.weight", + "diffusion_model.blocks.0.mlp.up.lora_up.weight", + "diffusion_model.blocks.0.mlp.down.lora_up.weight"]) + self.assertEqual(self.det(s), "krea2") + + def test_kohya_underscore_form(self): + s = _sd(["lora_unet_blocks_0_attn_wq.lora_down.weight", + "lora_unet_blocks_0_mlp_gate.lora_up.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", + "diffusion_model.blocks.3.mlp.gate.lora_up.weight"]) + self.assertEqual(self.det(s), "krea2") + + def test_no_collision_with_wan_flux_qwen(self): + # Krea is checked before WAN; must not steal these, and they must not steal Krea. + wan = _sd(["diffusion_model.blocks.0.self_attn.q.a", + "diffusion_model.blocks.0.cross_attn.k.b", + "diffusion_model.blocks.0.ffn.0.c"]) + flux = _sd(["diffusion_model.double_blocks.0.img_attn.qkv.lora_up.weight"]) + qwen = _sd(["transformer.transformer_blocks.0.attn.to_q.a", + "transformer.transformer_blocks.0.img_mlp.net.a"]) + self.assertEqual(self.det(wan), "wan") + self.assertEqual(self.det(flux), "flux") + self.assertEqual(self.det(qwen), "qwen_image") + + @unittest.skipIf(torch is None, "torch is not installed") class AnimaNormalizationTests(unittest.TestCase): """All trainer forms normalize to canonical diffusion_model.blocks.N.* keys."""