feat(Model Support): add Krea-2-Turbo/Raw model + LoRA support - #9304
Conversation
Integrate Krea-2-Turbo (krea/Krea-2-Turbo) text-to-image per
NEW_MODEL_INTEGRATION.md: Krea2Transformer2DModel (single-stream MMDiT)
+ Qwen3-VL text encoder (12-layer hidden-state tap, 4D prompt_embeds)
+ reused Qwen-Image VAE + FlowMatchEulerDiscrete scheduler.
Backend:
- taxonomy: BaseModelType.Krea2, ModelType/ModelFormat.Qwen3VLEncoder,
Krea2VariantType (Turbo = "krea2_turbo" to avoid Z-Image collision)
- config probes: Main_Diffusers/Checkpoint_Krea2, Qwen3VLEncoder,
LoRA_LyCORIS_Krea2 (text_fusion/time_mod_proj signature; excluded
from the Qwen-Image probe to avoid double-match)
- loaders for the diffusers pipeline + standalone Qwen3-VL encoder,
with runtime workarounds for the HF model's version mismatches
(AutoTokenizer, extra_special_tokens={}, rope_parameters->rope_scaling)
- native sampling (pack/unpack, position_ids, linear-mu shift) and
hand-written Euler denoise loop; reuses qwen_image l2i/i2l
- invocations: model_loader, text_encoder, denoise, lora_loader, plus
two ecosystem enhancers (conditioning rebalance, seed variance)
- LoRA conversion for diffusers PEFT (lora_transformer- prefix)
Frontend:
- 'krea-2' base + qwen3_vl_encoder type/format across model maps,
buildKrea2Graph, addKrea2LoRAs, graph-builder denoise/base lists,
optimal dimension 1024, regenerated schema.ts
Fixes:
- estimate transformer working memory in krea2_denoise so the cache
reserves activation headroom and offloads more model under partial
loading; fixes fp8 + LoRA OOM at 1024 (model was placed before LoRA
patches were applied, leaving no room for their activations)
WIP: requires diffusers main (>=0.39 dev) for Krea2Transformer2DModel;
pyproject.toml temporarily pins diffusers to git main.
|
Amazing! I was just thinking of working on this myself and you did it for me! |
Allow non-diffusers Krea-2 transformers (GGUF/fp8) to run with standalone single-file VAE + Qwen3-VL encoder, fixing several blockers found in testing. - buildKrea2Graph: drop the hard "requires Diffusers-format" assert; instead require both a VAE and a Qwen3-VL encoder to be selected when the transformer is not diffusers (mirrors readiness.ts). - Qwen3-VL encoder remap: handle both single-file key conventions — implicit (model.layers.*) and explicit (model.language_model.*). The old blind model.* -> language_model.* turned the bf16 file's keys into language_model.language_model.* (398 meta tensors -> "Cannot copy out of meta tensor" crash). Both files now load 0 missing / 0 unexpected / 0 meta. - Qwen3-VL tokenizer/config: broaden the offline-cache fallback from OSError to Exception so a partial HF cache (config present, vocab missing) re-fetches instead of dying with TypeError. - Qwen3-VL encoder fp8: keep an fp8 source checkpoint fp8-resident with per-layer upcast (storage float8_e4m3fn, compute bf16) instead of dequantizing to bf16. Halves resident VRAM (~8.9GB -> ~4.4GB), avoiding partial-load thrashing alongside a large transformer. Auto-enabled for fp8 sources on CUDA; bf16 files stay bf16. - Qwen-Image VAE: a native-layout qwen_image_vae single file is classified with the Anima base and loaded as AutoencoderKLWan, but the qwen l2i/i2l nodes need AutoencoderKLQwenImage. Add backend/krea2/vae_compat.py::as_qwen_image_vae to reinterpret a Wan VAE as AutoencoderKLQwenImage (state dicts are identical, 194/194 keys); both qwen VAE nodes use it. Idempotent for real QwenImage VAEs.
|
Can you add it to the list of supported models here https://github.com/invoke-ai/InvokeAI/blob/main/README.md?plain=1#L61 |
An upstream merge reintroduced the AutoencoderKLQwenImage isinstance asserts in the qwen VAE nodes (without the import → F821) and dropped the adapter in the i2l path. A native-layout qwen_image_vae single file is classified with the Anima base and loaded as AutoencoderKLWan, so the asserts fail at runtime. - qwen_image_latents_to_image: drop the reintroduced pre-device assert (the as_qwen_image_vae adapter inside model_on_device already handles the class). - qwen_image_image_to_latents: restore the as_qwen_image_vae import + adapter call, remove both asserts. - estimate_vae_working_memory_qwen_image only reads tensor shape + element size, so it runs correctly on either VAE class before the adapter.
…l_loader
krea2_model_loader was added to MainModelLoaderNodes but not to
isMainModelWithoutUnet, and the guard wasn't a type predicate — so it never
narrowed modelLoader. OutputFields of the loader union collapses to the common
'vae' field, making g.addEdge(modelLoader, 'unet', ...) in addInpaint/addOutpaint
fail to type-check ('unet' not assignable to 'vae').
Redefine the guard as a type predicate keyed on the inverse (only
main_model_loader/sdxl_model_loader expose a unet), so every transformer-based
loader is treated as unet-less automatically and the negated branch narrows to
the unet-bearing loaders.
zKrea2VariantType was only used within common.ts (in zAnyModelVariant) and never referenced externally, so knip flagged it as an unused export. Every sibling variant enum avoids this by being asserted in common.test-d.ts; add the missing Krea2VariantType assertion, which both uses the export and verifies the manual zod enum matches the generated S['Krea2VariantType'].
…layout Follow-up to the seed-variance recalibration: the metadata recall parsers still used the old ranges, so recalling an image dispatched state the backend rejects. - Krea2SeedVarianceStrength recall now parses 0..2 (the std-multiplier range), not 0..100 — recalling the old absolute value 20 no longer produces invalid state that buildKrea2Graph forwards to a failing generation. - Krea2SeedVarianceRandomizePercent recall now allows 0 (the disabled value), matching the slider, param state, and invocation. - LoRA_LyCORIS_Krea2_Config accepts a transformer-only LoRA using the diffusion_model.transformer_blocks.* layout under an explicit Krea-2 override; the converter already handles the diffusion_model. prefix. Adds range boundary tests for both recall parsers and diffusion_model.* LoRA accept/reject tests.
…pport Resolves conflicts with upstream video generation (invoke-ai#9163), Ideogram 4 (invoke-ai#9303), T5 GGUF encoder (invoke-ai#9324) and the Qwen VAE device fix (invoke-ai#9373). Notable resolutions: - qwen_image_latents_to_image: keep the as_qwen_image_vae() reinterpretation but adopt upstream's vae_info.compute_device fix (invoke-ai#9373) - graphBuilderUtils: keep the allow-list isMainModelWithoutUnet predicate, which covers wan_model_loader automatically - generationSettingsVisibility: add 'wan' and 'ideogram-4' to
Brings in the newer Krea-2 fixes (seed variance calibration, metadata recall ranges, diffusion_model LoRA layout) on top of the local merge of upstream video generation (invoke-ai#9163). Conflict resolutions: - Took origin's ordering/formatting for main.py, factory.py and qwen_image_latents_to_image.py (content-identical) - Kept the local side wherever Wan/video code is involved: starter model bundles, BASES_WITHOUT_STANDARD_SCHEDULER, and the isWan guards plus ParamWanModelSelects block in AdvancedSettingsAccordion
JPPhoto
left a comment
There was a problem hiding this comment.
Not too many to fix, hopefully:
-
invokeai/backend/model_manager/configs/lora.py:918_has_complete_lora_pair_for_prefixes()accepts a file after finding any complete pair. A file containing one valid layer plus an orphanlora_A,lora_B,lora_down, orlora_upinstalls successfully, then fails during LoRA conversion at generation time. Automatic detection likewise checks only for any suffix. Validate every LoRA layer pair during installation. Test: accept multiple complete pairs; reject each mixed complete-plus-orphan case. -
invokeai/frontend/web/src/features/parameters/components/Krea2Enhancers/ParamKrea2RebalanceWeights.tsx:14dispatches arbitrary text, readiness does not validate it, andbuildKrea2Graph()forwards it to the backend. Wrong counts, nonnumeric values,nan, orinflet generation queue beforeKrea2ConditioningRebalanceInvocation._parse_weights()fails.Krea2RebalanceWeights.parse()also recalls any string into this invalid state. Test: accept exactly 12 finite numbers; block or reject wrong counts, nonnumeric values,nan, andinffrom both UI state and metadata recall. -
invokeai/backend/krea2/vae_compat.py:19:as_qwen_image_vae()accepts everyAutoencoderKLWan, including 48-channel, patchified Wan 2.2 VAEs. Directqwen_image_i2lorqwen_image_l2iworkflows now pass the former class guard, then fail on 16-vs-48-channel normalization or produce incompatible latents. Require the Qwen-compatible 16-channel, 8x-spatial configuration. Test: accept default 16-channel Wan/Qwen-compatible models; reject 48-channel and patchified configurations before encode/decode.
…ompatible VAE Three install/queue-time guards so malformed inputs are rejected up front instead of failing mid-generation: - LoRA identification now requires every lora_A/B (or lora_down/up) weight to have its partner half. A valid layer plus a dangling half previously installed and then crashed during LoRA conversion; both the explicit-override and the automatic-detection paths now validate completeness. - Krea-2 Conditioning Rebalance weights are validated as exactly 12 finite numbers before generation: in readiness (blocks the queue), in metadata recall (rejects instead of dispatching invalid state), and in the input field (isInvalid). Mirrors Krea2ConditioningRebalanceInvocation._parse_weights. - as_qwen_image_vae now requires the Qwen-Image geometry (16 latent channels, 8x spatial, no patchification) and rejects Wan 2.2's 48-channel / patchified VAE before encode/decode, rather than failing on 16-vs-48 normalization. Adds LoRA orphan-pair tests, rebalance-weight validator + recall tests, and Wan VAE geometry accept/reject tests.
JPPhoto
left a comment
There was a problem hiding this comment.
New list:
-
invokeai/backend/model_manager/configs/lora.py:943: explicit Krea-2 override validates only approved prefixes. A validtransformer_blockspair plus an orphantext_fusionhalf installs, then conversion raises. Test: reject any orphaned half anywhere in the converted state dict. -
invokeai/backend/model_manager/configs/lora.py:997:dora_scalealone satisfies LoRA detection, while pair validation accepts files containing no A/B or down/up weights. Loading later fails on missing weights. Test: rejecttext_fusionfiles containing onlydora_scale; accept complete A/B plusdora_scale. -
invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts:1097: JavaScriptNumberaccepts hexadecimal, binary, and octal values such as0x10; backendinvokeai/app/invocations/krea2_conditioning_rebalance.py:49uses Pythonfloat, which rejects them. UI permits graphs guaranteed to fail. Test: reject hex, binary, and octal tokens while accepting decimal and scientific notation. -
invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py:109: standard PEFT DoRA keylora_magnitude_vectoris unrecognized, split into a bogus layer, then rejected despite complete A/B weights. Valid Diffusers DoRA adapters cannot load. Test: convert A/B pluslora_magnitude_vector; expectDoRALayerpreserving magnitude. -
invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py:49:transformer.anddiffusion_model.aliases normalize to the same final key; line 80 silently overwrites the first layer. Mixed-layout adapters lose weights based on dictionary order. Test: provide both aliases for one logical layer with different tensors; require explicit rejection instead of overwrite.
Addresses five install/convert-time issues so malformed Krea-2 LoRAs and rebalance weights are rejected up front (or converted correctly): - Explicit Krea-2 override now rejects an orphaned lora_A/B (or lora_down/up) half anywhere in the state dict, not just under the approved prefixes — a transformer_blocks pair plus a dangling text_fusion half previously installed and then crashed during conversion. - Krea-2 LoRA detection now requires a complete weight pair; a file with only dora_scale (no A/B weights) is rejected instead of failing later on load. - Rebalance weights are restricted to decimal/scientific notation, rejecting the hex/binary/octal literals (0x10, 0b10, 0o10) that JS Number() accepts but the backend's Python float() rejects. - The converter now recognizes the standard PEFT/Diffusers DoRA magnitude key lora_magnitude_vector.weight, mapping it to dora_scale so a valid DoRA adapter loads as a DoRALayer instead of being split into a bogus layer. - Conflicting transformer./diffusion_model. aliases that normalize to the same target layer now raise explicitly instead of silently overwriting one. Adds tests for each case.
…ction
Native Krea-2 LoRAs (e.g. sliders) name modules differently from InvokeAI's
diffusers Krea2Transformer2DModel: diffusion_model.blocks.N with attn.wq/wk/wv/
wo/gate, mlp.{down,gate,up}, and a txtfusion stage. These were misidentified as
Anima (whose strict detector matched the bare blocks.N.mlp.*) and, even when
forced to Krea-2, could not be applied because the converter only understood the
diffusers PEFT layout.
- Add a verified 1:1 native->diffusers key remap in the Krea-2 LoRA converter
(blocks->transformer_blocks, attn.wq/wk/wv->to_q/to_k/to_v, attn.wo->to_out.0,
attn.gate->to_gate, mlp->ff, txtfusion->text_fusion). Every native module maps
onto a real Linear in the diffusers model (checked against all 512 keys of a
real slider LoRA). DoRA magnitude survives the remap.
- Extend Krea-2 LoRA detection (config + converter) to recognize the native
signature (txtfusion, or the gated attention attn.wq + attn.gate).
- Tighten the Anima strict detector to require the Anima-specific mlp.layer_N /
mlp_layerN naming instead of a bare mlp, so a native Krea-2 LoRA is no longer
false-matched as Anima. No Anima/Wan regressions.
Adds native remap, DoRA-through-remap, diffusers-untouched, and native
identification tests.
…chensack/InvokeAI into feat/krea2-turbo-support
|
Just fixed:
|
Krea-2's transformer uses grouped-query attention (48 query / 12 KV heads) and its stock processor calls scaled_dot_product_attention with enable_gqa=True. PyTorch only supports enable_gqa on the math SDPA backend, which materializes the full O(seq^2) score matrix: ~6.75 GB per attention at 1280x720 (3600 tokens) and ~40 GB at 2560x1440. On builds without flash attention (e.g. Windows) there is no fused fallback, so generation either OOMs or the model cache offloads the transformer to RAM and the forward pass crawls. - Add Krea2MemoryEfficientAttnProcessor: expands the KV heads (repeat_interleave) so enable_gqa is not needed, and runs under the memory-efficient SDPA kernel (O(seq) memory, supports the padding mask). Numerically equivalent to the stock processor; measured ~6.75 GB -> ~1.41 GB per block at 3600 tokens. Installed on the transformer in krea2_denoise before the denoise loop. - Recalibrate _estimate_working_memory: with O(seq) attention the activation footprint is small and ~linear, so the previous ~2.6 MiB/token (O(seq^2)) figure no longer applies. The new estimate reserves realistic headroom (~8.5 GB at 2560x1440 instead of an impossible ~36 GB), so the idle Qwen3-VL encoder is evicted and the fp8 transformer stays resident on a 24 GB card. Adds processor equivalence tests (GQA and non-GQA) and a working-memory bound regression test.
…chensack/InvokeAI into feat/krea2-turbo-support
JPPhoto
left a comment
There was a problem hiding this comment.
Approved as-is, hopefully you can do final testing. This is appropriate for a follow-up PR:
invokeai/backend/model_manager/load/model_loaders/krea2.py:_convert_krea2_native_to_diffusersandinvokeai/backend/model_manager/load/model_loaders/krea2.py:_remap_qwen3vl_singlefile_keys: Both normalizers silently overwrite keys when native and normalized aliases target the same tensor, making malformed mixed-layout checkpoints order-dependent. Test: Supply distinct tensors under colliding aliases in both orders and assert explicit rejection.
…erwriting Both Krea-2 key normalizers (native->diffusers transformer keys, ComfyUI single-file Qwen3-VL encoder keys) mapped each source key to one target key and wrote it straight into the output dict. A malformed mixed-layout checkpoint that carries both a native key and its already-normalized alias (e.g. blocks.0.attn.wq.weight and transformer_blocks.0.attn.to_q.weight, or a bare layers.1.weight and its model.-prefixed twin) collapses both onto one target key, and the surviving tensor depended on dict iteration order. Route every write through a shared _put_unique_key helper that raises an actionable RuntimeError naming both colliding source keys, so such a checkpoint fails at load time instead of silently dropping a tensor. Add order-independent collision regression tests for both normalizers.
|
@Pfannkuchensack I found two more items but they should absolutely be in a follow-up so we don't thrash on this PR:
|
Summary
Integrate Krea-2 text-to-image per
NEW_MODEL_INTEGRATION.md: both Krea-2-Turbo (krea/Krea-2-Turbo, distilled) and Krea-2-Raw (krea/Krea-2-Raw, undistilled Base). Architecture:Krea2Transformer2DModel(single-stream MMDiT, ~12B) + Qwen3-VL text encoder (12-layer hidden-state tap → 4Dprompt_embeds) + reused Qwen-Image VAE (AutoencoderKLQwenImage) +FlowMatchEulerDiscreteScheduler.is_distilled=true): fixedmu=1.15, 8 steps, CFG off (cfg 1.0).is_distilled=false): resolution-awaremu, ~28 steps, CFG ~4.5. Variant is read from the pipelineis_distilledflag (single-file/GGUF fall back to a filename heuristic).Formats: full Diffusers pipeline, single-file checkpoint (incl. ComfyUI scaled fp8), and GGUF (Q2–Q8). Single-file/GGUF ship only the transformer, so a standalone Qwen-Image VAE + Qwen3-VL encoder are selected in the loader UI (enforced by readiness before enqueue). NVFP4 intentionally skipped (needs Blackwell FP4 kernels).
VRAM: fp8 layerwise-cast weight storage for the transformer (diffusers + single-file) and for the fp8 Qwen3-VL encoder (~8.9 GB bf16 → ~4.4 GB resident), keeping 1024² + LoRA within 24 GB.
diffusers dependency (resolved)
Krea2Transformer2DModel/Krea2Pipelinelanded in diffusers 0.39.0 (stable).pyproject.tomlnow pinsdiffusers[torch]==0.39.0anduv.lockis updated accordingly — the previous git-main blocker is gone.Notable extras
isMainModelWithoutUnetis now a proper type predicate covering all transformer-based loaders).QA Instructions
uv sync --extra cuda(pulls diffusers 0.39.0); confirmpython -c "from diffusers import Krea2Transformer2DModel, Krea2Pipeline".Krea-2-Turboand aKrea-2-Rawfolder. Confirm they probe asmain / diffusers / krea-2with variantkrea2_turbo/krea2_base, and the bundled encoder asqwen3_vl_encoder. On model select, Turbo defaults to 8 steps / cfg 1.0, Raw to ~28 steps / cfg 4.5.vantagewithai/Krea-2-Turbo-GGUF) + the standalone Qwen-Image VAE + Qwen3-VL encoder. Confirm the loader UI requires the VAE + encoder before enqueue and that generation succeeds. The fp8 encoder logsFP8 layerwise casting enabled for Qwen3-VL encoder.lora.lycoris.krea-2(not qwen-image), applies, and that 1024² + LoRA + fp8 does not OOM.Tests
_has_krea2_keys, GGUF/checkpoint/diffusers variant, default settings) and the single-file Qwen3-VL encoder probe (visual-tower vs. text-only Qwen3). Run:pytest tests/backend/model_manager/configs/test_krea2_main_config.py tests/backend/model_manager/configs/test_qwen3_vl_encoder_config.py.Merge Plan
diffusers blocker is resolved (pinned to stable 0.39.0). No DB schema changes.
paramsSlicegains Krea-2 fields with a corresponding migration.Out of scope / follow-ups
Checklist
What's Newcopy (if doing a release after this PR)