feat(Model Support): add Krea-2-Turbo model + LoRA support (WIP)

[Pfannkuchensack]

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 → 4D prompt_embeds) + reused Qwen-Image VAE (AutoencoderKLQwenImage) + FlowMatchEulerDiscreteScheduler.

• Turbo (is_distilled=true): fixed mu=1.15, 8 steps, CFG off (cfg 1.0).
• Raw (is_distilled=false): resolution-aware mu, ~28 steps, CFG ~4.5. Variant is read from the pipeline is_distilled flag (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 / Krea2Pipeline landed in diffusers 0.39.0 (stable). pyproject.toml now pins diffusers[torch]==0.39.0 and uv.lock is updated accordingly — the previous git-main blocker is gone.

Notable extras

• Conditioning enhancers (Advanced Options, under CFG Scale, default OFF, with tooltips): Conditioning Rebalance (per-layer weighting) and Seed Variance Enhancer (restores per-seed diversity on the distilled model).
• Metadata recall for the standalone VAE, Qwen3-VL encoder, and both enhancers, so single-file/GGUF generations reproduce from metadata.
• Starter models: Krea-2 Turbo, Krea-2 Raw, Krea-2 Turbo GGUF (Q4_K_M / Q8_0, with VAE + encoder dependencies), and a standalone Qwen3-VL 4B encoder.
• Fixed a latent type-guard bug surfaced by the new loader union (isMainModelWithoutUnet is now a proper type predicate covering all transformer-based loaders).

QA Instructions

1. Dependency: uv sync --extra cuda (pulls diffusers 0.39.0); confirm python -c "from diffusers import Krea2Transformer2DModel, Krea2Pipeline".
2. Install (Diffusers): point InvokeAI at a Krea-2-Turbo and a Krea-2-Raw folder. Confirm they probe as main / diffusers / krea-2 with variant krea2_turbo / krea2_base, and the bundled encoder as qwen3_vl_encoder. On model select, Turbo defaults to 8 steps / cfg 1.0, Raw to ~28 steps / cfg 4.5.
3. Generate (txt2img): 1024², enable FP8 in the model's Default Settings on 24 GB cards. Confirm an image renders for both Turbo and Raw.
4. Single-file / GGUF: install a GGUF transformer (e.g. 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 logs FP8 layerwise casting enabled for Qwen3-VL encoder.
5. LoRA: load a Krea-2 LoRA (diffusers PEFT). Confirm it probes as lora.lycoris.krea-2 (not qwen-image), applies, and that 1024² + LoRA + fp8 does not OOM.
6. CFG: at cfg 1.0 (Turbo) confirm no negative prompt is run and recall shows no negative prompt; at cfg > 1 (Raw) confirm negative conditioning is used.
7. Enhancers: with both off, output matches stock. Enable Conditioning Rebalance → stronger prompt adherence; enable Seed Variance → meaningfully different images across seeds. Confirm tooltips render and that recall restores both.
8. img2img / inpaint / outpaint: round-trip an init image.

Tests

• Config-probe unit tests for Krea-2 variant detection (name heuristic, _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. paramsSlice gains Krea-2 fields with a corresponding migration.

Out of scope / follow-ups

• ControlNet / IP-Adapter plumbing exists but cannot be validated until a Krea-2 control checkpoint is released.

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)

[lstein]

Amazing! I was just thinking of working on this myself and you did it for me!

[kappacommit]

Can you add it to the list of supported models here https://github.com/invoke-ai/InvokeAI/blob/main/README.md?plain=1#L61

[JPPhoto]

@Pfannkuchensack Here's a list of things that should probably to be addressed:

• pyproject.toml:39 requires diffusers[torch]==0.39.0, but uv.lock:632 and uv.lock:1128 still resolve 0.37.0. Because docker/Dockerfile:77 runs uv sync --frozen, Docker installations retain a Diffusers release without the new Krea2 classes. Regenerate and commit uv.lock.

• invokeai/app/invocations/krea2_text_encoder.py:69 ignores self.qwen3_vl_encoder.loras. Meanwhile, invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py:54 converts text-encoder patches and invokeai/frontend/web/src/features/nodes/util/graph/generation/addKrea2LoRAs.ts:35 routes them to the encoder. Qwen3-VL portions of Krea2 LoRAs are silently ignored. Apply them with LayerPatcher and KREA2_LORA_QWEN3VL_PREFIX, following z_image_text_encoder.py. To expose this issue, add a test that supplies a LoRA containing a text-encoder patch and asserts that the encoder layer is patched during _encode().

• invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts:106 handles component cleanup and default selection for Z-Image, Anima, FLUX.2, and Qwen Image, but not Krea2. Standalone Krea2 VAE and encoder selections survive switches to other model families. When a Krea2 Diffusers model is later selected, invokeai/frontend/web/src/features/nodes/util/graph/generation/buildKrea2Graph.ts:65 passes those stale selections as overrides instead of using the bundled components. To expose this issue, add a listener test that selects Krea2 GGUF components, switches away, then selects Krea2 Diffusers and asserts that the stale component selections are cleared.

• The same missing Krea2 branch in invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts means installed GGUF starter dependencies are not selected automatically. Generation remains blocked by invokeai/frontend/web/src/features/queue/store/readiness.ts:327 until the user manually selects the VAE and encoder. To expose this issue, add a listener test with installed Qwen Image VAE and Qwen3-VL encoder records and assert that selecting a Krea2 GGUF model populates both fields.

• invokeai/frontend/web/src/features/metadata/parsing.tsx:1162 defines recall handlers for the Krea2 VAE, Qwen3-VL encoder, seed variance settings, and rebalance settings, but invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataActions.tsx:32 omits every one of them. These values are absent from the Recall Parameters tab and cannot be recalled individually. To expose this issue, extend ImageMetadataActions.test.tsx to require every Krea2 metadata handler in IMAGE_METADATA_ACTION_HANDLERS.

• Krea2SeedVarianceEnabled.parse() at invokeai/frontend/web/src/features/metadata/parsing.tsx:1213 and Krea2RebalanceEnabled.parse() at invokeai/frontend/web/src/features/metadata/parsing.tsx:1269 return false when their fields are absent and do not gate that fallback on the current model being Krea2. Recalling metadata from an unrelated or older image can therefore clear hidden Krea2 enhancer state. To expose this issue, add a metadata recall test that starts with the Krea2 settings enabled, recalls non-Krea2 metadata without these keys, and checks whether those settings should remain unchanged.

• invokeai/backend/model_manager/starter_models.py:1932 does not add a Krea2 entry to STARTER_BUNDLES. The launchpad renders this mapping, so the newly defined Krea2 models and their dependencies are not offered as a Krea2 launchpad bundle. To expose this issue, add a starter-model response test asserting that BaseModelType.Krea2 is present and that its GGUF option includes the Qwen Image VAE and Qwen3-VL encoder dependencies.

• The new user-facing model family, approximately 26 GB Diffusers installation, standalone component requirements, GGUF workflow, Raw versus Turbo behavior, and enhancer controls are undocumented. Add Krea2 hardware requirements to docs/src/content/docs/start-here/system-requirements.mdx and an appropriate model usage page under docs/src/content/docs/features/.

• Krea2GGUFCheckpointModel._load_from_gguf() in invokeai/backend/model_manager/load/model_loaders/krea2.py:365 detects parameters left on the meta device but only logs a warning and returns the incomplete model. Unsupported or partially converted GGUF layouts will fail later during inference instead of being rejected during loading. To expose this issue, add a loader test with a required key missing and assert that loading raises a descriptive exception.

• Krea2CheckpointModel._load_from_singlefile() in invokeai/backend/model_manager/load/model_loaders/krea2.py:326 uses strict=False and discards the missing and unexpected key lists. A misidentified, incomplete, or differently converted checkpoint may retain meta parameters or silently omit weights. To expose this issue, add a checkpoint loader test with missing required transformer weights and assert that the loader rejects it before returning the model.

• Qwen3VLEncoderCheckpointLoader._load_text_encoder() in invokeai/backend/model_manager/load/model_loaders/krea2.py:535 also uses strict=False without inspecting missing keys or checking for remaining meta parameters. A partially compatible Qwen3-VL checkpoint can be accepted and fail only when prompt encoding begins. To expose this issue, add a loader test with an incomplete language-model tower and assert that loading fails with the missing key names.

• _has_krea2_lora_keys() in invokeai/backend/model_manager/configs/lora.py:904 requires a LoRA to target text_fusion or time_mod_proj. A legitimate Krea2 LoRA trained only against shared transformer_blocks cannot be installed as Krea2 and may instead be classified as Qwen Image. Explicitly selecting Krea2 as an override does not bypass the same validation. To expose this issue, add a config-probe test for a transformer-block-only Krea2 LoRA and define how ambiguous LoRAs should be classified or manually overridden.

• Krea2DenoiseInvocation._prepare_cfg_scale() in invokeai/app/invocations/krea2_denoise.py:124 requires a list length equal to the number of timesteps after img2img schedule clipping. A custom workflow supplying one CFG value per configured step will fail when denoising_start or denoising_end reduces the active steps. To expose this issue, add a denoise test with eight CFG values, eight configured steps, and a nonzero denoising_start.

• denoising_start and denoising_end are independently constrained in invokeai/app/invocations/krea2_denoise.py:72, but nothing requires start to be less than or equal to end. A reversed range can produce an empty sigma slice, after which sigmas_sched[0] is accessed before the zero-step guard. To expose this issue, add an invocation validation test using denoising_start=0.9 and denoising_end=0.1 and require a clear validation error.

• Supplying denoise_mask without initial latents reaches assert init_latents is not None in invokeai/app/invocations/krea2_denoise.py:223. API and workflow callers can construct that graph, producing an assertion failure instead of a user-facing validation error. To expose this issue, add an invocation test with a mask but no initial latents.

• Krea2DenoiseInvocation._run_diffusion() in invokeai/app/invocations/krea2_denoise.py:155 hardcodes torch.bfloat16, while the Krea2 model loaders use TorchDevice.choose_bfloat16_safe_dtype(). On a device where the loader selects a fallback dtype, conditioning and inference still force bfloat16 and may fail or trigger unsupported operations. To expose this issue, mock a device whose safe dtype is float32 or float16 and assert that denoising consistently uses that selected dtype.

• Krea2TextEncoderInvocation._encode() in invokeai/app/invocations/krea2_text_encoder.py:114 also hardcodes torch.bfloat16, independently of the dtype selected when the Qwen3-VL model was loaded. Use the loaded model's effective dtype or the same safe-dtype helper. To expose this issue, load or mock an encoder using a non-bfloat16 dtype and assert that its output conditioning uses a compatible dtype.

• _is_distilled() in invokeai/app/invocations/krea2_denoise.py:132 catches every exception and defaults to Turbo behavior. Configuration lookup failures, malformed model_index.json, and filesystem errors are silently converted into a fixed Turbo shift, which can produce incorrect sampling for a Raw model. To expose this issue, make context.models.get_config() fail and assert that the invocation reports the configuration error instead of silently changing variants.

• Main_Diffusers_Krea2_Config._get_variant() in invokeai/backend/model_manager/configs/main.py also catches every exception and defaults to Turbo. A malformed or unreadable Raw model index can therefore be registered with Turbo defaults. To expose this issue, add a model-probe test with malformed model_index.json and require an identification error rather than a Turbo classification.

• Krea2LoRALoaderInvocation in invokeai/app/invocations/krea2_lora_loader.py:54 checks only that the model exists; unlike the collection loader, it does not verify that the selected LoRA has BaseModelType.Krea2. Direct API or workflow callers can attach an unrelated LoRA and receive an apparently valid output whose patches are ignored or mismatched. To expose this issue, add an invocation test passing a FLUX or SDXL LoRA and require a compatibility error.

• Krea2LoRACollectionLoader in invokeai/app/invocations/krea2_lora_loader.py:137 deduplicates only entries within the incoming collection. It does not check LoRAs already present on the supplied transformer or encoder, so chaining collection loaders can apply the same LoRA more than once. To expose this issue, start with a field containing one LoRA, pass the same LoRA in the collection, and assert that it is rejected or retained only once.

• as_qwen_image_vae() in invokeai/backend/krea2/vae_compat.py:20 accepts any model that is not already AutoencoderKLQwenImage and attempts to reinterpret its entire state dictionary. The modified Qwen Image encode and decode nodes therefore no longer reject an incompatible VAE immediately; they instantiate a large replacement module before strict loading fails. Add an explicit AutoencoderKLWan type check. To expose this issue, connect an unrelated VAE and assert that it is rejected before constructing a Qwen Image VAE.

• The Krea2 model-loader and denoise paths have no coverage for full Diffusers loading, single-file conversion, GGUF conversion, Qwen3-VL prompt encoding, LoRA application, schedule clipping, CFG, img2img, inpaint, outpaint, cancellation, or model-cache cleanup. Add focused coverage under tests/backend/model_manager/load/ and tests/app/invocations/, plus graph construction coverage alongside invokeai/frontend/web/src/features/nodes/util/graph/generation/buildKrea2Graph.ts.

• The new diffusers upgrade affects every existing model family, but the PR changes only the dependency declaration and Krea2 code. Add representative compatibility coverage for existing FLUX, Qwen Image, Z-Image, SD, and Anima loader paths against Diffusers 0.39.0.

• The numeric placeholder in invokeai/frontend/web/src/features/parameters/components/Krea2Enhancers/ParamKrea2RebalanceWeights.tsx:26 is a direct user-visible string rather than an en.json translation key. Move it into invokeai/frontend/web/public/locales/en.json to keep all newly introduced UI text within the localization system.

[JPPhoto]

I fixed a few of the light effort issues:

• Updated uv.lock to Diffusers 0.39.0 with minimal lockfile churn.
• Added input validation in invokeai/app/invocations/krea2_denoise.py for empty/inverted denoising ranges and masks without initial latents.
• Restricted as_qwen_image_vae() in invokeai/backend/krea2/vae_compat.py to compatible VAE classes.

[JPPhoto]

Got these in as well:

• Added all Krea handlers to the metadata recall action list.
• Prevented the generic VAE handler from duplicating the dedicated Krea VAE action.
• Gated Krea scalar metadata using the image's model base, preventing unrelated images from exposing or recalling Krea defaults.
• Localized the rebalance-weight placeholder through en.json.

[JPPhoto]

Current state. I'll work on these:

• invokeai/backend/model_manager/load/model_loaders/krea2.py:409 loads standalone directory-format Qwen3-VL encoders without copying rope_parameters into rope_scaling. The bundled Diffusers and single-file paths explicitly perform this compatibility fix because the current Transformers implementation can otherwise crash. Apply the same normalization here. Add a standalone-directory loader test where rope_scaling is absent and rope_parameters is present.

• invokeai/backend/patches/lora_conversions/krea2_lora_conversion_utils.py:85 discards a layer's stored .alpha and .dora_scale whenever the layer uses PEFT lora_A/lora_B keys. This changes explicit LoRA scaling and turns DoRA weights into ordinary LoRA weights. Preserve both values during normalization. Add conversion tests asserting that per-layer alpha is retained and that a PEFT DoRA layer produces a DoRA patch with its magnitude tensor.

• invokeai/app/invocations/krea2_lora_loader.py:69 and invokeai/app/invocations/krea2_lora_loader.py:122 assume the transformer and encoder LoRA lists are synchronized. If a LoRA is already attached to only one field, the single loader raises before repairing the other field, while the collection loader's global added_loras list silently skips both. The result can be a transformer patch without its Qwen3-VL patch, or the reverse. Deduplicate independently per output field. Test both asymmetric starting states against both loader invocations.

• invokeai/backend/model_manager/configs/lora.py:904 still requires text_fusion or time_mod_proj keys to classify a LoRA as Krea-2. A legitimate transformer-block-only Krea-2 LoRA is indistinguishable from Qwen Image automatically, but it also cannot be installed through an explicit Krea-2 selection because the same probe rejects it. Provide a deliberate override path for ambiguous LoRAs. Test automatic ambiguity plus successful explicit Krea-2 installation for transformer.transformer_blocks.*.lora_A/B weights.

• invokeai/app/invocations/krea2_denoise.py:196 enables CFG globally when any entry in a list exceeds 1.0, and invokeai/app/invocations/krea2_denoise.py:354 then applies the CFG formula at every active step. A schedule such as [4.0, 1.0, 0.5] therefore applies unconditional blending at the 1.0 and 0.5 steps even though the invocation documents values <= 1 as disabling CFG. Decide CFG per step, using the conditional prediction directly when that step's scale is <= 1. Test mixed schedules, including clipped schedules whose active window contains no value above 1.0.

• invokeai/app/invocations/krea2_denoise.py:245 rounds fractional denoising bounds independently. A valid narrow interval can round to identical indices, and invokeai/app/invocations/krea2_denoise.py:278 silently returns undenoised latents. Without input latents this returns raw noise for subsequent VAE decoding; with input latents it silently performs a no-op or noise blend. Reject intervals that produce zero effective steps, or explicitly define safe zero-step behavior. Test low step counts with bounds such as 0.0 to 0.01.

• invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts:326 and invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts:580 treat a missing model-config cache entry as a non-Diffusers Krea-2 model. During startup, metadata recall, or direct dispatch, selecting a Diffusers model before getModelConfigs is populated can auto-select standalone components that buildKrea2Graph.ts later passes as stale overrides. The listener does not rerun when the cache arrives. Defer component changes until the format is known, or obtain the selected model's individual cached configuration. Test selection with an initially absent list cache followed by a Diffusers configuration response.

• invokeai/backend/model_manager/configs/factory.py:424 now considers any root config.json an unambiguous model marker and bypasses the general-purpose-directory file-count guard. config.json is generic, so a large application or data directory can now be hashed and registered as an unknown model instead of being rejected. Restrict the bypass to recognized model configuration content or model_index.json. Test a large non-model directory containing an unrelated config.json.

• invokeai/backend/model_manager/configs/qwen3_vl_encoder.py:46 classifies a standalone Qwen3-VL directory solely from its config class. It does not require encoder weights or tokenizer assets, although Qwen3VLEncoderLoader uses local_files_only=True for both. An incomplete directory can therefore install successfully, satisfy readiness, and fail only during generation. Validate the required weight index/files and tokenizer files for direct and nested layouts. Add rejection tests for config-only, weights-only, and tokenizer-only directories.

• The current focused tests do not exercise actual full Diffusers Krea-2 loading, standalone Qwen3-VL directory loading, single-file or GGUF model construction, prompt encoding with LoRA patches, CFG numerical behavior, or complete img2img/inpaint/outpaint execution. Add small fixture-based loader and invocation tests that reach model construction and forward boundaries rather than stopping at state-dict utilities and graph shape assertions.

• The upgrade to diffusers==0.39.0 changes a shared dependency used by existing FLUX, Qwen Image, Z-Image, SD, and Anima loaders, but this PR contains no representative compatibility coverage for those paths. Add smoke coverage that constructs or loads one representative configuration from each affected family against the pinned version.