feat(krea2): training-free style reference - #9512
Conversation
…nodes Pulls two seams out of the Krea-2 nodes so out-of-tree node packs can add per-token behaviour without copying large chunks of core. krea2_denoise: the attention wiring inside _run_diffusion (state creation, processor installation, per-pass mask install, cleanup) becomes four overridable methods. Subclasses can install their own processors and a richer state object without duplicating the 281-line _run_diffusion. krea2_text_encoder: the encode body moves to invokeai/backend/krea2/text_encoding.py, so the exact token layout (prefix drop, post-truncation suffix concat, mRoPE positions) is stated once. It takes an optional build_token_values callback which receives the body's offset mapping and returns one value per token; the result is extended over the suffix and sliced by the same prefix drop as the embeddings, so a caller's per-token vector stays aligned with the conditioning by construction. Pure refactor: the generated OpenAPI schema is byte-identical to before, no node versions change, and without the callback the tokenizer receives the exact same call it always did.
github-actions
Bot
added
python
There was a problem hiding this comment.
One thing to fix:
invokeai/backend/krea2/text_encoding.py:121-122: suffix_values is always CPU, so a valid non-CPU callback result fails at torch.cat. Test: return a meta or CUDA tensor from build_token_values; current code raises a device-mismatch RuntimeError.
The exact torch.cat failure is a corner case; it requires the callback to return a non-CPU tensor. InvokeAI never passes the callback, so default generation is unaffected...
But, for the intended GPU extension use, CPU results are also hazardous: offset mappings are CPU tensors, so callbacks will naturally return CPU values; combining them with GPU embeddings can fail later.
Fix this by moving values to device or creating suffix values with body_values.new_ones(...).
Port nkxx188/ComfyUI-Krea2-StyleTransfer (MIT) as a core feature: transfer a
reference image's look via shared-KV reference attention, with no adapter model
and no extra weights.
Runs as two passes rather than upstream's doubled batch — the reference goes
through the transformer alone and its image-token K/V are stashed, then spliced
into the target pass. This is equivalent (upstream's reference rows never attend
to the target) and avoids re-padding the text conditioning, which krea2_denoise
strips rather than masks.
K/V are captured before the GQA head expansion (12 kv heads, not 48): provably
identical, and 4x smaller — 0.5 GiB at 1024x1024 instead of 2.0 GiB. The cache is
declared to the model cache and cleared via the exit stack so it cannot leak onto
the cached transformer. At 2560x1440 the combined footprint no longer fits a 24 GB
card alongside the model; this is asserted in a test rather than hidden.
Backend adds style_reference{,_rf,_extension}.py plus a krea2_style_reference
node; krea2_denoise gains style_reference and an optional reference prompt.
Frontend wires Krea-2 into the existing Reference Images panel following the Wan
2.2 pattern (image, no adapter model) with a style strength slider.
…kuchensack/InvokeAI into feat/krea2_style_reference
Style reference had added a third positional parameter to _install_attention_processors. Custom node packs override that seam with the documented two-argument form, so the change broke every generation they ran — style reference connected or not. Style reference now installs itself in _install_style_reference_processors, which re-installs the processors on top of whatever the seam produced and refuses outright when the seam has been overridden, rather than silently discarding a subclass's processors. Also raise the style-reference activation margin from 1.2x to 1.35x: an end-to-end run at 1024x1024 measured a ~1.6 GiB peak-VRAM increase, of which ~0.5 GiB is the captured K/V cache, so the arithmetic-derived 20% was short. Under-estimating here makes the model cache offload the transformer to RAM, which looks like a hang rather than an error.
Krea-2's style reference splices a single reference's attention keys/values into the target, so the graph builder consumes exactly one image. The Reference Images panel happily accepts several, and the rest were dropped without a word — which reads as "all of them are being used". Every usable Krea-2 reference image after the first now carries a warning, shown on both the entity header and its preview thumbnail. It is advisory: generation still proceeds with the first image. Deliberately implemented as a separate, sibling-aware validator rather than folded into getGlobalReferenceImageWarnings — the graph builder filters its candidates on that function returning no warnings, so folding it in would have excluded the one image we do use.
JPPhoto
left a comment
JPPhoto
left a comment
There was a problem hiding this comment.
Merge blockers:
invokeai/frontend/web/src/features/nodes/util/graph/generation/buildKrea2Graph.ts:203-212: style node copiesdenoise.width/heightbefore mode helpers assign them (addTextToImage.ts:40-55, etc.). JSON omitsundefined; backend defaults to 1024x1024, then rejects normal non-1024 generation. Effect: style-reference generation fails. Likelihood: normal supported resolutions. Recovery: manual graph edit/retry. Test: build actual txt2img/img2img graph at 768x1024; assert style node dimensions equal denoise dimensions. Existing test mocks sizing, so both values remainundefined.
Other findings/issues:
-
invokeai/backend/krea2/style_reference.py:100-120,invokeai/app/invocations/krea2_denoise.py:423-543:style_strength=0claims to disable reference, but still VAE-encodes it, captures K/V, runs one extra transformer pass per step, and retains VRAM cache. Effect: “off” still costs roughly 2x runtime and may trigger VRAM offload/OOM. Likelihood: normal slider use. Recovery: disable/delete reference manually. Test: run with strength 0 and assert no style capture/pass/cache. -
invokeai/backend/krea2/text_encoding.py:117-122: callback may return CUDA/non-CPU values, but suffix values are always CPU;torch.catthen raises device-mismatch. Effect: intended out-of-tree encoder extension crashes. Likelihood: rare; optional callback has no in-tree caller, but this PR explicitly exposes it for extension packs. Recovery: patch callback/upgrade PR. Test: callback returnstorch.ones(..., device="meta")or CUDA tensor. -
invokeai/frontend/web/src/features/nodes/util/graph/generation/buildKrea2Graph.ts:194-216: new user-facing style-reference feature has no documentation, while PR metadata still claims “pure refactor”, “backend-only”, and “no functional change”. Effect: reviewers/users miss one-image, dimension, runtime, and VRAM constraints. Likelihood: current PR. Recovery: update PR body anddocs/src/content/docs/features/krea-2.mdx. Test: compare current head against documented QA/body claims.
Suggestions:
- Instead of creating style node before mode sizing, assign dimensions first or patch style node after the mode helper.
- Instead of running style machinery at strength 0, omit the style node or short-circuit denoise.
- Instead of CPU
torch.ones(...), usebody_values.new_ones(...).
Addresses review 4977240290 on invoke-ai#9512. The style-reference node was built before the generation-mode helper ran, so it copied `denoise.width`/`height` while they were still undefined. JSON drops undefined, the backend fell back to its 1024x1024 default, and the denoise node then rejected the reference at every other resolution. The block now runs after the mode helper, with an assert as a backstop. The test did not catch it because the mode helpers were mocked as no-ops, leaving both sides undefined; the mocks now assign the dimensions the way the real helpers do, and the sizing is asserted at 1024x1024, 768x1024 and 1152x896 across all four generation modes. A style strength of 0 is documented as disabling the reference, but the machinery ran regardless: a VAE encode, a capture pass per step and a retained K/V cache, all for an attention mix of 0 — roughly double the runtime and enough extra VRAM to trigger offloading. The graph builder now omits the node entirely at 0, and the denoise node treats a strength of 0 like an unconnected input, so the workflow-editor path is covered too. The next enabled reference takes over, exactly as it does for a disabled one, and the "only one reference image is used" validator filters the same way so its warning stays on the right entity. Field descriptions and the UI comment are corrected to match. `build_token_values` is an extension seam, so it may hand back a tensor on the encoder's device, but the suffix was always built on the CPU and the concat would raise. `new_ones` inherits both device and dtype. Documents the style reference in docs/src/content/docs/features/krea-2.mdx: what it is, how to reach it from the canvas and the workflow editor, and the four constraints — one reference image, the reference must match the denoise resolution, ~2x runtime, and 0.5 GB at 1024² / 1.7 GB at 1440p.
Pfannkuchensack
changed the title
3 participants
Summary
Adds training-free style reference for Krea-2 — transfer the look of a reference image (palette,
texture, rendering) while the prompt keeps driving the content. There is no adapter model and no LoRA:
the reference's attention keys and values are spliced into the target's, so it works with any Krea-2
checkpoint out of the box.
The PR started as the refactor described at the bottom, and that refactor is what the feature is built
on. It is not a pure refactor any more: it adds nodes, fields, a canvas control and a redux config type.
Backend
krea2_style_referencenode (Style Reference - Krea-2, v1.0.0) — VAE-encodes the reference at agiven width/height and emits the new
Krea2StyleReferenceField.fitchooses how the reference'saspect ratio is reconciled with the target (
crop/contain/stretch).krea2_denoise(v1.2.0 → v1.3.0) gains two optional connection-only inputs,style_referenceand
style_reference_conditioning. The latter lets a short neutral prompt drive the reference pass;left unconnected it reuses the positive conditioning at no extra cost.
invokeai/backend/krea2/style_reference.pyholds the attention-side math,style_reference_rf.pythereference noising schedule,
style_reference_extension.pythe capture/inject lifecycle. The denoisenode stays free of style bookkeeping.
Per step the reference runs one extra transformer pass whose image-token K/V are stashed, then the
target's passes (conditional and unconditional) run inside the injection. Styling only the conditional
pass would make CFG amplify
styled_cond - plain_uncond, which overshoots badly above cfg ~4, so thesingle reference pass is shared by both. The retained K/V is included in the working-memory estimate so
the model cache offloads for it instead of OOMing.
Frontend
krea2_reference_imageref-image config with a single Style Strength slider (0–2, default 1). Nomodel picker — there is no adapter model to choose.
model-specific ref-image types.
being dropped silently.
Constraints
target's and share its rotary embedding. The canvas wires this up; in the workflow editor, set the same
width/height on both nodes or the denoise node refuses the reference with a message naming both sides.
24 GB card alongside the model.
node ignores a connected reference at 0, so it costs nothing.
The underlying refactor
Two extension seams were extracted from the Krea-2 nodes so out-of-tree node packs can add per-token
behaviour without copying large parts of core. The style reference is the first consumer of the first
seam.
krea2_denoise: overridable attention wiring. The attention setup inside_run_diffusionwas inlineand the processor factory it calls is a module-level function, so a subclass could only intervene by
duplicating the whole 281-line method. It is now four small methods:
_install_attention_processors(transformer, exit_stack)Krea2RegionalPromptingState, installs the memory-efficient processors, registers cleanup_clear_attention_state(state)set_attention_mask(None)_build_attention_payload(extension, dtype)extension.get_attention_mask(), computed once before the loop_install_attention_payload(state, payload)Style reference installs itself through the same seam and does not widen its signature — a subclass
that overrides it keeps working, and one that installs its own processors is rejected with a clear error
rather than silently losing its styling.
krea2_text_encoder: encoding moved to the backend. The encode body moves toinvokeai/backend/krea2/text_encoding.py, so Krea-2's token layout — the prefix drop, thepost-truncation suffix concat, the mRoPE position ids — is stated in one place. It gains one optional
parameter,
build_token_values: a callback receiving the tokenizer's offset mapping for the prompt bodyand returning one value per token. The result is extended over the suffix and sliced by the same prefix
drop as the embeddings, so a caller's per-token vector stays aligned by construction. Callers that don't
need it leave it unset and the tokenizer call is unchanged.
Related Issues / Discussions
None.
QA Instructions
Style reference, canvas:
the style node is built after the generation-mode helper assigns the denoise dimensions, so it is
encoded at the same size.
krea2_style_referencenode and therun takes the same time as with no reference at all.
Style reference, workflow editor:
Style Reference - Krea-2→Denoise - Krea-2with mismatched width/height must fail with amessage naming both sizes, not a shape error inside attention.
Style Reference Promptwith a short neutral prompt and confirm less of the reference'ssubject bleeds into the result.
Refactor, unchanged behaviour:
return_offsets_mappingis only passed when a callback is supplied. Covered by a test.restricts alternating main blocks only.
Tests cover: the capture/inject order and that both CFG passes are styled; the size and shape rejections;
cache cleanup when the transformer raises; the strength-0 bypass on both sides; the working-memory
estimate including the K/V cache; the graph builder's sizing across all four generation modes and several
resolutions; the one-reference warning; and each extension seam's default plus a subclass replacing it.
Merge Plan
Nothing special. No DB migration. The new redux config type is added to the existing
RefImageStateunion, which is versioned and validated by zod — no slice migration needed.
Checklist
docs/src/content/docs/features/krea-2.mdxWhat's Newcopy (if doing a release after this PR)