Krea 2: Support multiple conditionings

[JPPhoto]

Summary

Adds support for multiple global positive and negative conditioning inputs on the Denoise - Krea-2 node.

• Accepts either one Krea2ConditioningField or a collection.
• Concatenates embeddings and attention masks along the text-token dimension.
• Preserves existing single-conditioning workflows.
• Rejects empty conditioning collections with a clear error.
• Bumps krea2_denoise from version 1.0.0 to 1.1.0.
• Documents collection behavior and the lack of spatial-mask support.

Related Issues / Discussions

QA Instructions

1. Create two Prompt - Krea-2 nodes.
2. Collect both conditioning outputs.
3. Connect the collection to Denoise - Krea-2 positive conditioning.
4. Run generation and confirm denoising succeeds.
5. Repeat with a negative conditioning collection using Krea-2 Raw and CFG greater than 1.0.
6. Confirm existing workflows with one conditioning still run unchanged.

Merge Plan

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[JPPhoto]

2 prompts, keeping the subject prompt but altering the style prompt to yield different results:

[lstein]

Adversarial review of the diff: I traced the new conditioning state through the transformer, checked the frontend field-type contract, ran the test suite, and probed the empty-collection paths directly. CI is fully green.

The approach is right and matches the flux/anima precedent. One item I'd like addressed before merge (working-memory accounting), plus a cheap change that resolves it and pays for itself.

What I attacked and found clean

• Mask dtype regression on the existing single-conditioning path. Every masked conditioning now routes through mask.to(dtype=torch.bool) (krea2_denoise.py:164), which it didn't before. krea2_text_encoder.py:145 already emits .bool(), so it's a no-op; downstream transformer_krea2.py:498 uses encoder_attention_mask.new_ones(...), so bool propagates consistently into the joint mask. No regression.
• 4D embeds. Krea2ConditioningInfo.prompt_embeds is (B, seq, 12, hidden) — unlike every other model family. dim=1 is the token axis (correct), and the fallback torch.ones(embeds.shape[:2]) yields (B, seq) (correct) despite the extra axis.
• Rotary positions. prepare_position_ids puts all text rows at (0,0,0), and text_seq_len is derived from the concatenated embeds (krea2_denoise.py:349), with the negative pass keeping its own ids. Concatenation can't desync the rotary length. As a bonus this makes the result genuinely invariant to collection order.
• Frontend contract. anyOf: [ref, array-of-ref, null] → parseFieldType.ts filters the null, sees 2 entries, returns SINGLE_OR_COLLECTION. The new "title" keys in the openapi diff match every other single-or-collection conditioning field — cosmetic.
• Version bump. Same major, so getMayUpdateNode offers an in-place update for saved 1.0.0 workflows rather than breaking them.

1. _estimate_working_memory ignores text tokens, which this PR makes scale with N

krea2_denoise.py:391-395 sizes the cache reservation purely from image_seq_len, and the docstring at :477 states the fixed 1.5 GB base covers "the resolution-independent overhead (… the text-fusion stage …)". Text is no longer resolution-independent — it is now N × 512.

Every Krea-2 conditioning is exactly 512 tokens regardless of prompt length: krea2_text_encoder.py:106-113 pads to body_max_length = 512 + 34 − 5 = 541, appends 5 suffix tokens, then drops 34. So four conditionings put 2048 text tokens into a stage sized for 512. Krea2TextFusion.forward reshapes to (B*seq, 12, dim) and runs attention + SwiGLU over that, so its footprint is strictly linear in N — at N=4 it plausibly exceeds the entire 1.5 GB base on its own.

This is exactly the failure the estimator was added to prevent, and the documented symptom is the bad one: the cache offloads the transformer to RAM and generation effectively hangs rather than raising a clean OOM. Please thread the post-concat text_seq_len (for both the cond and uncond sequences) into _estimate_working_memory.

2. Each extra conditioning costs a full 512 tokens, overwhelmingly padding

Following from the fixed padding above: two 8-word prompts cost 1024 text tokens of which ~1000 are masked-out padding. At 1024×1024 (4096 image tokens) that is +11% sequence and roughly +23% attention FLOPs per step for nothing; at N=4, +33% sequence and roughly +78% FLOPs.

Since padding is already excluded as attention keys, dropping masked tokens before torch.cat is numerically identical for every surviving token — softmax over valid keys is unchanged, text rotary positions are all zero, and the output slice hidden_states[:, text_seq_len:] adapts. That makes N-prompt cost proportional to actual prompt length instead of N×512, largely defuses finding 1, and in the common case leaves an all-valid mask that can be returned as None — which additionally lets SDPA use the flash backend instead of falling back to EFFICIENT_ATTENTION (attention.py:21).

3. An empty negative collection behaves inconsistently

has_negative_conditioning = self.negative_conditioning is not None (krea2_denoise.py:319) treats [] as present. I confirmed both branches by running them:

• cfg_scale=4.0, negative_conditioning=[] → hard failure: ValueError: At least one Krea-2 conditioning is required.
• cfg_scale=1.0, negative_conditioning=[] → runs fine, the empty collection is silently ignored.

An empty collection is reachable from the editor (a Collect node with nothing wired into it). Since the input is explicitly optional, [] most naturally means "no negative conditioning", and treating it as absent would be consistent across both branches. At minimum the message should name the offending input — as written it doesn't say whether the positive or negative side was empty.

4. Duplicated chat-template markers land mid-sequence

Each conditioning carries its own <|im_end|>\n<|im_start|>assistant\n suffix, so the concatenated sequence contains assistant-turn markers in the middle — a token layout Krea-2 never saw during training. Not wrong, but not quite "one longer prompt" either, and it tends to surface as subtle quality loss rather than an error. Given the docs section claims the conditionings are "combined into one longer text token sequence", worth a QA pass comparing [promptA, promptB] against a single concatenated prompt string.

5. Test gap: single conditioning with a mask is uncovered

The three new tests all pass lists. The pre-existing runtime tests do exercise the single-field path end-to-end, but _runtime_context builds its conditionings with prompt_embeds_mask=None — so the single-field-plus-mask case, which is what every real workflow hits today and which now flows through the new torch.cat + bool-cast branch, has no coverage. A one-line test asserting that a lone masked field round-trips its mask unchanged would lock that in.

[JPPhoto]

@lstein Thanks for the review. Addressed the actionable feedback in 4966e8596e:

• Removed masked padding before concatenating conditionings.
• Includde the longest positive/negative text sequence in the working-memory estimate.
• Treated an empty optional negative collection as absent.
• Added coverage for single masked conditioning, multiple compacted conditionings, empty negatives, and text-length memory estimates.
• Clarified the independently encoded concatenation behavior in the Krea-2 docs.

Each conditioning still retains its independently encoded Krea chat wrapper and assistant suffix. The denoise node receives embeddings rather than raw prompts, so joining prompts before encoding would require a different contract. This is not expected to be equivalent to encoding one combined prompt; comparing their relative image quality remains an empirical QA test.

[JPPhoto]

Note: I tested outputs compared to the prior ones that were not stripped and they're almost identical.

[lstein]

Re-reviewed 4966e8596e. All five items are addressed, and I verified the load-bearing claim numerically rather than by inspection. CI is green across all 17 checks.

Verification of the padding-strip equivalence

This was the change with the most room to be subtly wrong, so I built a tiny Krea2Transformer2DModel (2 blocks, 1 layerwise + 1 refiner fusion block) and compared the two paths on the same inputs: (A) padded embeds + concatenated key-padding mask, versus (B) stripped embeds + encoder_attention_mask=None, using two conditionings with non-contiguous masks ([T,T,F,F,T] and [T,F,T,F,F], so 10 padded slots collapse to 5):

[stock attn processor]  padded_seq=10 stripped_seq=5  max|A-B|=1.19e-07  rel=8.4e-08
[Krea2MemoryEfficient]  padded_seq=10 stripped_seq=5  max|A-B|=1.79e-07  rel=1.3e-07

Float32 epsilon, under both the stock processor and Krea2MemoryEfficientAttnProcessor. That confirms it end-to-end: the mask only ever acted as a key-padding mask in Krea2TextFusion.refiner_blocks and the joint blocks (transformer_krea2.py:494-499), the layerwise blocks and SwiGLU/RMSNorm are per-token, and text rotary rows are all (0,0,0) — so deleting padded rows is exactly equivalent for every surviving token. Your "almost identical" observation on real images is the expected bf16 kernel/reduction-order difference, not a semantic one. Passing encoder_attention_mask=None is correctly handled (transformer_krea2.py:492-493), and it now unlocks the flash SDPA backend as a bonus.

Item-by-item

1. Working memory — resolved. (image_seq_len + max(pos, neg)) * per_token_bytes is the right shape: the passes are sequential so max rather than sum is correct, and per-text-token cost in the fusion stage works out to roughly 12 × 6912 × 2 × 2 ≈ 330 KB, comfortably inside the 0.5 MB/token constant — so dropping the text-fusion stage from the fixed-base docstring is justified rather than just moved around.
2. Padding removal — resolved, see above. Cost is now proportional to real prompt length instead of N × 512.
3. Empty negative collection — resolved. I checked the truthiness switch rather than assuming: Krea2ConditioningField defines neither __bool__ nor __len__, so bool(field) is True and only None/[] fall through — a single wired negative conditioning cannot be silently dropped. As a side effect the empty-positive error message is now unambiguous, since the negative side can no longer reach it.
4. Chat-template markers — agreed, and your reasoning is right: the node consumes embeddings, so pre-encode joining would need a different contract. The docs wording ("independently encoded conditionings are concatenated after padding tokens are removed") describes what actually happens. Leaving it as an empirical QA question is the correct call.
5. Test coverage — resolved by test_load_text_conditioning_compacts_a_single_masked_conditioning, which is exactly the real-world path. 38 tests pass locally.

Attacks that came up clean

• Existing producers vs. the new shape check. krea2_conditioning_rebalance.py:69-75 and krea2_seed_variance.py:59-82 both forward prompt_embeds_mask alongside rewritten embeds — both are strictly elementwise, so mask.shape == embeds.shape[:2] holds and neither can trip the new ValueError.
• Mixed masked/unmasked conditionings. A mask-less conditioning keeps its padding and now has no mask to exclude it — but that was already true pre-PR (encoder_attention_mask=None) and in the first revision (torch.ones fallback), so no regression. krea2_text_encoder.py always emits a mask anyway.
• Batch > 1 with unequal valid counts now raises where it previously worked. Unreachable: the encoder always produces B=1.
• All-False mask compacts to a zero-length text sequence rather than erroring. Also unreachable — the 5 suffix tokens are always unmasked, so every real conditioning has at least 5 valid tokens.
• Non-contiguous masks. Boolean indexing preserves order, and the equivalence harness above used non-contiguous masks specifically to check this.

One heads-up rather than a request: a single-conditioning workflow now compacts 512 tokens down to the true prompt length, so identical seeds shift by float epsilon against 1.0.0 output. You've already spot-checked that. Also purely cosmetic: _load_text_conditioning now always returns None as its second element, so the tuple[Tensor, Tensor | None] annotation and both *_prompt_mask locals are vestigial — harmless, only worth a mention if you'd rather not leave a mask channel that can never carry a value.

Approving.