feat: Support Krea2 - #14589
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR adds full support for the Krea 2 (K2) single-stream MMDiT diffusion model. A new 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
comfy/model_base.py (1)
2286-2291: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winUse
CONDCrossAttnforc_crossattnto keep mixed-length prompt batching efficient.Wrapping
cross_attnwithCONDRegularcan disable cross-attention concat for differing token lengths and force extra UNet passes.🔧 Proposed fix
def extra_conds(self, **kwargs): out = super().extra_conds(**kwargs) cross_attn = kwargs.get("cross_attn", None) if cross_attn is not None: - out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn) + out['c_crossattn'] = comfy.conds.CONDCrossAttn(cross_attn) return outAs per path instructions,
comfy/**reviews should guard hot-path performance regressions.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy/model_base.py` around lines 2286 - 2291, In the extra_conds method, change the wrapper type for the c_crossattn assignment from CONDRegular to CONDCrossAttn when wrapping the cross_attn parameter. The line that currently reads out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn) should use CONDCrossAttn instead to maintain efficient mixed-length prompt batching and avoid unnecessary UNet passes.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy/model_detection.py`:
- Around line 837-850: The Krea2 detection block is not inferring txtheads and
txtkvheads configuration values, leaving them to default in SingleStreamDiT
which causes shape mismatches for non-default checkpoints. Add two new lines to
the dit_config dictionary to infer txtheads and txtkvheads by extracting and
processing the attention weights from the txtfusion blocks (similar to how heads
and kvheads are already being extracted from blocks.0.attn.wq.weight and
blocks.0.attn.wk.weight). Use the same head_dim value and reference the
corresponding txtfusion attention weight keys in the state_dict to divide by
head_dim, ensuring the configuration matches the actual model structure.
In `@comfy/sd.py`:
- Around line 1603-1605: The elif block handling CLIPType.KREA2 (lines
1603-1605) is incorrectly placed in the QWEN3_4B branch, but Krea2 is only
compatible with Qwen3-VL models. Remove the entire elif clause that checks for
CLIPType.KREA2 and sets clip_target.clip and clip_target.tokenizer with the
krea2 text encoder, as this functionality should only be present in the
QWEN3VL_* branch where it is properly supported.
---
Nitpick comments:
In `@comfy/model_base.py`:
- Around line 2286-2291: In the extra_conds method, change the wrapper type for
the c_crossattn assignment from CONDRegular to CONDCrossAttn when wrapping the
cross_attn parameter. The line that currently reads out['c_crossattn'] =
comfy.conds.CONDRegular(cross_attn) should use CONDCrossAttn instead to maintain
efficient mixed-length prompt batching and avoid unnecessary UNet passes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: cf1c7b46-2428-4682-b356-db4e2869e196
📒 Files selected for processing (8)
comfy/ldm/krea2/model.pycomfy/lora.pycomfy/model_base.pycomfy/model_detection.pycomfy/sd.pycomfy/supported_models.pycomfy/text_encoders/krea2.pynodes.py
works but poorly, rather not allow
|
@kijai Did you look at this ? https://www.reddit.com/r/StableDiffusion/comments/1ud2nyq/comment/ot8mwm6/?utm_source=share&utm_medium=web3x&utm_name=web3xcss&utm_term=1&utm_content=share_button It seems there's a slight issue with this implementation |
The implementation is fine, the warning is a small easter egg from the model creators. go figure it out before the official release tomorrow ;) |
|
style-reference system seems missing |
They did not release this. |
5 participants
Add Krea2 model support
Adds support for the Krea2 image model.
comfy/ldm/krea2/model.py:SingleStreamDiTtransformer (single-stream blocks, text-fusion transformer, SwiGLU MLP, separate Q/K/V/O attention).model_base.py,model_detection.py,supported_models.py, andsd.py; new CLIP type intext_encoders/krea2.py.lora.pymaps diffusers-format LoRA keys (transformer.*, lycoris) onto Krea2's original weight names so diffusers-trained LoRAs apply directly.