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 <noreply@anthropic.com>
This commit is contained in:
Ethanfel
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
f073ff87c1
commit
26828757eb
+17
-3
@@ -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":
|
||||
|
||||
+1
-1
@@ -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",
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user