[Pfannkuchensack]

Summary

The Krea-2 "ControlNets" people ask for (pose, depth) are ai-toolkit edit LoRAs: they only work when the prompt is encoded together with reference images by the Qwen3-VL vision tower and the reference latents ride along in the transformer at timestep 0. Invoke dropped the vision tower and had no reference path, so these LoRAs did nothing useful.

This adds Krea-2's edit mode to the backend and the workflow editor:

• Prompt - Krea-2 gains reference_images (up to 3). They are read by the Qwen3-VL vision tower, which is now its own model-cache entry (SubModelType.VisionEncoder) loaded only when images are given and lent to the text encoder for that forward. Text-only encoding is unchanged and loads nothing new. Supported encoders: the Krea-2 Diffusers model's own encoder and Qwen3-VL single files/folders that include the vision weights; GGUF encoders are refused with a clear message.
• Denoise - Krea-2 gains reference_latents (up to 3): reference tokens appended after the target with their own position ids, modulated at t=0, output limited to the target tokens. A reference count that differs from the prompt's is refused; combining references with Style Reference or regional prompts is refused.
• Reference K/V cache (reference_kv_cache, off by default) for LoRAs trained with ai-toolkit's kv_cache option (references attend only to each other): the reference K/V are computed once per run at t=0 and appended in every pass. It is also faster.
• Image to Latents - Qwen Image gains a reference mode with ai-toolkit's Krea-2 reference sizing (about 1 MP, sides rounded to 16, bilinear, RGB).

Related Issues / Discussions

Follows the request for Krea-2 pose/depth ControlNet support; the Generate/Canvas UI follows in a separate PR.

QA Instructions

RTX 4090, Krea-2 Turbo (int8 single file) + qwen3vl_4b_fp8_scaled, 10 steps, CFG 1, 1024², LoRA weight 1.0:

Run	Result
cicalooo/krea2_unidepth_depth, depth map as reference, joint mode, 3 seeds	follows the depth map (head, shoulders, waterline) on every seed
thedeoxen/Krea-2-pose-controlnet, DWPose skeleton, joint mode	blotches and extra figures — this LoRA was trained with isolated references
same LoRA, Reference K/V cache on, 3 seeds + a second pose	clean, follows the pose and matches the LoRA's published examples
references without a LoRA	the model copies the reference into the image, as expected

Timing per step: 1 reference 2.8 s joint vs 1.5 s cached; 3 references 6.5 s joint vs ~2.0 s cached. Peak VRAM 9.7 GB with 3 references.

Numerical parity on real weights: the transformer forward with references is torch.equal to the Ostris/ComfyUI formulation (Turbo Q4_K_M); the encoder is torch.equal to ai-toolkit's encode_krea_prompt layout when both get the same pixels (CPU vs GPU bicubic resizing alone moves bf16 image tokens by a similar amount as ±2.4e-5 pixel noise does). On CPU: the split block equals the stock block, the no-reference forward equals the stock forward, the cached path equals a joint pass with ai-toolkit's isolating mask, lending the tower equals a model that never lost it.

env -u HF_ENDPOINT uv run --no-sync pytest -n 4 tests/backend/krea2 tests/backend/model_manager tests/backend/qwen3_vl tests/backend/architectures tests/model_identification tests/app: 11714 passed; the one failure (test_model_cache_vram_accounting) measured real free VRAM while another process was using the GPU and passes on its own. Ruff clean.

Review

Resolved: the rebalance and seed-variance nodes now carry the reference count (they would have reset it and blocked every edit graph); clear errors for encoders without vision weights and for references the processor rejects; the negative prompt's count is checked only when CFG is above 1; the K/V-cache attention-score estimate prices the real (non-square) peak; docs and field descriptions corrected. Decided: references drop alpha (as in training) instead of compositing over white; the tower loads before the encoder so a GGUF encoder is refused before the language tower is read. Remaining: the Krea-2 facet's max_reference_images stays at the single style reference until the UI PR; the encoder working-memory estimate reuses Qwen3-VL-8B measurements; a LoRA file does not say whether it needs the K/V cache.

Compatibility / Rollout

New SubModelType.VisionEncoder value and new optional fields on krea2_text_encoder (1.2.0), krea2_denoise (1.4.0) and qwen_image_i2l (1.2.0); existing workflows keep their behaviour. Krea2ConditioningInfo gains num_reference_images (default 0, so stored conditionings load unchanged). openapi.json/schema.ts regenerated. No database migration.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[lstein]

Reviewed at 791d76f. No blockers; the items below are follow-ups, roughly most important first.

1. Vision-tower lock doesn't respect the encoder's working-memory reservation — invokeai/app/invocations/text_encoder/krea2_text_encoder.py:175
The encoder is locked with working_mem_bytes (:142), then the tower is locked with the default reserve (device_working_mem_gb). The tower's ~0.8 GB of weights can land in the headroom reserved for the encoder's forward, so on a card with little spare VRAM an edit-mode encode can OOM, or the tower partially loads. Passing the same working_mem_bytes to the tower lock, or adding the tower's size to the encoder's reservation, would fix it.

2. Edit-mode prompts are not length-capped — invokeai/backend/krea2/text_encoding.py, tokenize_krea2_prompt_with_images
The text-only path truncates the body to KREA2_BODY_MAX_LENGTH; the image path calls the processor without truncation. Encoder working memory grows with T², so a very long prompt asks for a very large reservation, and together with (1) it's the realistic way to hit the OOM. Is the missing cap intentional, i.e. does ai-toolkit not truncate here?

3. Negative-prompt reference count is checked against the unclipped CFG schedule — invokeai/app/invocations/krea2/krea2_denoise.py:406-409
_validate_reference_counts checks self.cfg_scale, but do_cfg (:547-553) uses the schedule clipped to [start_idx:end_idx]. Example: cfg_scale=[4,4,1,1,1,1,1,1], denoising_start=0.5, and a negative prompt encoded without references. The node refuses even though the negative prompt never runs. It's only over-strict and never produces wrong output.

4. webv1 zSubModelType is missing 'vision_encoder' — invokeai/frontend/webv1/src/features/nodes/types/common.ts:180
common.test-d.ts:51 asserts equality with the shared schema's SubModelType, which now includes vision_encoder, so the webv1 typecheck fails. CI doesn't run webv1, so this wouldn't show there. It's a one-line addition after 'processor', as was done for Qwen-Image-2.1.

5. Reference checks only compare counts

• _validate_reference_counts compares num_reference_images only. A prompt encoded with [depth, pose] and a denoise given [pose, depth], or two other images, is accepted silently, although the docs and error text tell users to keep the same order. Either say in the docs that the order isn't checked, or record identities in Krea2ConditioningInfo.
• Several unmasked positive prompts, each encoded with the same references, put the "Picture N" vision tokens into the sequence twice, which the LoRA wasn't trained on. Nothing refuses this or tests it. Consider refusing more than one conditioning when references are connected.

6. Smaller items

• A reference the processor rejects (an extreme aspect ratio) fails only after both the tower and the ~8 GB language tower have been loaded. Tokenizing with the processor before loading the encoder would fail fast.
• Only reference_kv_cache is refused for subclasses that override _install_attention_processors (krea2_denoise.py:386-394). In joint mode, a subclass whose processors are sized for [text | target] will get [text | target | refs].
• Docs (krea-2.mdx:184) say each reference is scaled to about 384×384, but prepare_vlm_reference only downscales; smaller images keep their size and the processor raises them to its pixel floor.
• Test gap: switching the vision-side resize in prepare_vlm_reference to antialias=False passes every test, because the tests check output size only, not pixel values.

Checked and found sound

• Lending the tower: detached in a finally, both entries locked for the whole window, no double counting of the cache size, LoRA patches applied before attach and removed after detach.
• K/V cache: equivalent to the joint pass with the isolating mask; sharing it between cond and uncond is correct; cleared on error and on exit; built after LoRA patching.
• Compatibility: old pickled Krea2ConditioningInfo loads with a count of 0, existing workflows and the non-reference qwen_image_i2l path are unchanged, and openapi.json regenerates identically.
• Tests: the new and changed tests pass on CPU (580 passed, 2 skipped). Most of the mutations I tried against the key lines were caught; the antialias change above was the one that got through.

Not checked: GPU runs, real Krea-2 Diffusers checkpoints, the VRAM figures, schema.ts regeneration.

[Pfannkuchensack]

Thanks for the review. Addressed in d9dfcd6:

1. Tower lock. The vision tower is now locked with the encoder's own working_mem_bytes. The cache already counts the locked encoder as allocated, so the tower only lands in VRAM beyond the encoder's reservation. Adding the tower's size to the encoder's reservation instead would count an already-resident tower twice. A test checks that both locks get the same reservation.
2. Edit prompts are not capped, on purpose. ai-toolkit and the Ostris edit pipeline tokenize the image path at its natural length, so the LoRAs are trained without truncation. The encoder's reservation uses the actual token count of that path, and with (1) it now covers the tower too.
3. Negative-prompt reference count. The check now uses the CFG schedule clipped to the denoise window, through the same helpers _run_diffusion uses for do_cfg. Your example ([4,4,1,1,1,1,1,1], denoising_start=0.5) is accepted, and the negative prompt is never loaded. Validation and the run read the same step count, and the window's rounding has its own test.
4. webv1. It isn't maintained or type-checked in CI, so I left it unchanged.
5. Reference checks.
   • Only the number of references is checked; the denoise only sees latents, so there is no identity to compare. The docs say so, and the error asks you to connect the same images in the same order.
   • A prompt collection of more than one prompt is now refused when reference latents are connected: always for the positive prompt, and for the negative when it runs. The error says to connect a single Prompt - Krea-2 node.
6. Smaller items.
   • The processor now tokenizes before the vision tower and the 8 GB language model load, so an image it rejects fails first. The GGUF refusal still comes before the language model, and the text-only path is unchanged.
   • A subclass that overrides _install_attention_processors is refused for any reference latents, not only the K/V cache.
   • The docs now say that larger references are scaled down to about 384×384 pixels of area, that smaller ones keep their size, that a side under 28 px is stretched to 28, and that the processor raises images below its pixel floor.
   • A pixel-value test pins the antialiased bicubic resize; antialias=False and a bilinear resize both fail it.

Checks: the Krea-2, Qwen3-VL tower and node-type tests pass on CPU (1528 passed), ruff is clean, and openapi.json is unchanged.

[lstein]

Re-reviewed d9dfcd67d6. Every item from the earlier review is addressed or explained, and I found no regressions. Approving.

Verified

• Vision-tower lock: the tower now locks with the encoder's working_mem_bytes. I agree that this is right and that adding the tower's size to the encoder's reservation would count it twice.
• Uncapped edit prompts: understood. They follow ai-toolkit, and the reservation uses the real token count.
• Negative-prompt check: validation and _run_diffusion now share _step_window and _active_cfg_scale. The window comes from self.steps, which is the count build_sigmas uses. Without denoising_start or denoising_end the behavior is unchanged.
• Other checks:
  • Prompt collections are refused when reference latents are connected.
  • A subclass that installs its own attention processors is refused for any reference latents.
  • The processor tokenizes before either tower loads, and the text-only path is unchanged.
• Docs: the sizing text and the note that only the count is checked match the code.
• Checks run:
  • ruff check and format on the changed files: clean.
  • Krea-2 denoise, text-encoder and backend tests: 259 passed, 2 skipped on CPU.
  • Mutations: tests caught each of these four reverted on its own: the tower reservation, the clipped CFG window, the refusal of several prompts and the refusal of subclasses.

Note

• webv1's zSubModelType still lacks vision_encoder, so the webv1 typecheck fails locally. CI doesn't run it, so this is fine as a known gap.

Not checked: GPU runs, real Krea-2 checkpoints and the VRAM figures.

#9923 is stacked on this PR, so this one should merge first.