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

README.md

@@ -75,6 +75,7 @@ Invoke features an organized gallery system for easily storing, accessing, and r
 - Flux.2 Klein 9B
 - Z-Image Turbo
 - Z-Image Base
+- Krea 2 Turbo
 - Anima
 - Qwen Image
 - Qwen Image Edit

invokeai/app/api/dependencies.py

@@ -60,6 +60,7 @@
     CogView4ConditioningInfo,
     ConditioningFieldData,
     FLUXConditioningInfo,
+    Krea2ConditioningInfo,
     QwenImageConditioningInfo,
     SD3ConditioningInfo,
     SDXLConditioningInfo,
@@ -153,6 +154,7 @@ def initialize(
                     CogView4ConditioningInfo,
                     ZImageConditioningInfo,
                     QwenImageConditioningInfo,
+                    Krea2ConditioningInfo,
                     AnimaConditioningInfo,
                 ],
                 ephemeral=True,

invokeai/app/invocations/fields.py

@@ -155,6 +155,7 @@ class FieldDescriptions:
     t5_encoder = "T5 tokenizer and text encoder"
     glm_encoder = "GLM (THUDM) tokenizer and text encoder"
     qwen3_encoder = "Qwen3 tokenizer and text encoder"
+    qwen3_vl_encoder = "Qwen3-VL tokenizer and text encoder"
     clip_embed_model = "CLIP Embed loader"
     clip_g_model = "CLIP-G Embed loader"
     unet = "UNet (scheduler, LoRAs)"
@@ -171,6 +172,7 @@ class FieldDescriptions:
     sd3_model = "SD3 model (MMDiTX) to load"
     cogview4_model = "CogView4 model (Transformer) to load"
     z_image_model = "Z-Image model (Transformer) to load"
+    krea2_model = "Krea-2 model (Transformer) to load"
     qwen_image_model = "Qwen Image Edit model (Transformer) to load"
     qwen_vl_encoder = "Qwen2.5-VL tokenizer, processor and text/vision encoder"
     sdxl_main_model = "SDXL Main model (UNet, VAE, CLIP1, CLIP2) to load"
@@ -349,6 +351,12 @@ class QwenImageConditioningField(BaseModel):
     conditioning_name: str = Field(description="The name of conditioning tensor")
 
 
+class Krea2ConditioningField(BaseModel):
+    """A Krea-2 conditioning tensor primitive value"""
+
+    conditioning_name: str = Field(description="The name of conditioning tensor")
+
+
 class AnimaConditioningField(BaseModel):
     """An Anima conditioning tensor primitive value.
 

invokeai/app/invocations/krea2_conditioning_rebalance.py

@@ -0,0 +1,74 @@
+import torch
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, Krea2ConditioningField
+from invokeai.app.invocations.primitives import Krea2ConditioningOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.krea2.sampling_utils import KREA2_SELECT_LAYERS
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
+    ConditioningFieldData,
+    Krea2ConditioningInfo,
+)
+
+_NUM_TEXT_LAYERS = len(KREA2_SELECT_LAYERS)  # 12
+
+
+@invocation(
+    "krea2_conditioning_rebalance",
+    title="Conditioning Rebalance - Krea-2",
+    tags=["conditioning", "krea2", "krea-2"],
+    category="conditioning",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class Krea2ConditioningRebalanceInvocation(BaseInvocation):
+    """Per-layer rebalancing of Krea-2 text conditioning (improves prompt adherence).
+
+    Krea-2 conditioning stacks 12 Qwen3-VL hidden-state layers per token. Weighting those layers
+    individually (and applying an overall multiplier) lets you push the model harder toward the prompt,
+    counteracting the quality-dilution from distillation. Ported from the ComfyUI
+    `ConditioningKrea2Rebalance` node. This is an optional pass between the text encoder and denoise.
+    """
+
+    conditioning: Krea2ConditioningField = InputField(
+        description=FieldDescriptions.cond, input=Input.Connection, title="Conditioning"
+    )
+    per_layer_weights: str = InputField(
+        default="1.0,1.0,1.0,1.0,1.0,1.0,1.0,2.5,5.0,1.1,4.0,1.0",
+        description=f"Comma-separated gains for the {_NUM_TEXT_LAYERS} tapped encoder layers (exactly "
+        f"{_NUM_TEXT_LAYERS} values).",
+    )
+    multiplier: float = InputField(
+        default=4.0,
+        description="Overall multiplier applied to the conditioning after per-layer weighting.",
+    )
+
+    def _parse_weights(self) -> list[float]:
+        try:
+            weights = [float(x.strip()) for x in self.per_layer_weights.split(",") if x.strip() != ""]
+        except ValueError as e:
+            raise ValueError(f"per_layer_weights must be comma-separated numbers: {e}") from e
+        if len(weights) != _NUM_TEXT_LAYERS:
+            raise ValueError(f"per_layer_weights must have exactly {_NUM_TEXT_LAYERS} values, got {len(weights)}.")
+        return weights
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> Krea2ConditioningOutput:
+        weights = self._parse_weights()
+
+        cond_data = context.conditioning.load(self.conditioning.conditioning_name)
+        assert len(cond_data.conditionings) == 1
+        conditioning = cond_data.conditionings[0]
+        assert isinstance(conditioning, Krea2ConditioningInfo)
+
+        embeds = conditioning.prompt_embeds  # (B, seq, 12, hidden)
+        gains = torch.tensor(weights, dtype=embeds.dtype, device=embeds.device).view(1, 1, _NUM_TEXT_LAYERS, 1)
+        embeds = embeds * gains * self.multiplier
+
+        new_data = ConditioningFieldData(
+            conditionings=[
+                Krea2ConditioningInfo(prompt_embeds=embeds, prompt_embeds_mask=conditioning.prompt_embeds_mask)
+            ]
+        )
+        conditioning_name = context.conditioning.save(new_data)
+        return Krea2ConditioningOutput.build(conditioning_name)