Feature: flux2 klein lora support

invokeai/app/invocations/flux2_denoise.py

@@ -38,6 +38,9 @@
 )
 from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType
 from invokeai.backend.patches.layer_patcher import LayerPatcher
+from invokeai.backend.patches.lora_conversions.flux_bfl_peft_lora_conversion_utils import (
+    convert_bfl_lora_patch_to_diffusers,
+)
 from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
@@ -503,11 +506,17 @@ def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor)
         return mask.expand_as(latents)
 
     def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]:
-        """Iterate over LoRA models to apply."""
+        """Iterate over LoRA models to apply.
+
+        Converts BFL-format LoRA keys to diffusers format if needed, since FLUX.2 Klein
+        uses Flux2Transformer2DModel (diffusers naming) but LoRAs may have been loaded
+        with BFL naming (e.g. when a Klein 4B LoRA is misidentified as FLUX.1).
+        """
         for lora in self.transformer.loras:
             lora_info = context.models.load(lora.lora)
             assert isinstance(lora_info.model, ModelPatchRaw)
-            yield (lora_info.model, lora.weight)
+            converted = convert_bfl_lora_patch_to_diffusers(lora_info.model)
+            yield (converted, lora.weight)
             del lora_info
 
     def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:

invokeai/app/invocations/flux2_klein_lora_loader.py

@@ -0,0 +1,182 @@
+"""FLUX.2 Klein LoRA Loader Invocation.
+
+Applies LoRA models to a FLUX.2 Klein transformer and/or Qwen3 text encoder.
+Unlike standard FLUX which uses CLIP+T5, Klein uses only Qwen3 for text encoding.
+"""
+
+from typing import Optional
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField
+from invokeai.app.invocations.model import LoRAField, ModelIdentifierField, Qwen3EncoderField, TransformerField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
+
+
+@invocation_output("flux2_klein_lora_loader_output")
+class Flux2KleinLoRALoaderOutput(BaseInvocationOutput):
+    """FLUX.2 Klein LoRA Loader Output"""
+
+    transformer: Optional[TransformerField] = OutputField(
+        default=None, description=FieldDescriptions.transformer, title="Transformer"
+    )
+    qwen3_encoder: Optional[Qwen3EncoderField] = OutputField(
+        default=None, description=FieldDescriptions.qwen3_encoder, title="Qwen3 Encoder"
+    )
+
+
+@invocation(
+    "flux2_klein_lora_loader",
+    title="Apply LoRA - Flux2 Klein",
+    tags=["lora", "model", "flux", "klein", "flux2"],
+    category="model",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class Flux2KleinLoRALoaderInvocation(BaseInvocation):
+    """Apply a LoRA model to a FLUX.2 Klein transformer and/or Qwen3 text encoder."""
+
+    lora: ModelIdentifierField = InputField(
+        description=FieldDescriptions.lora_model,
+        title="LoRA",
+        ui_model_base=BaseModelType.Flux2,
+        ui_model_type=ModelType.LoRA,
+    )
+    weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight)
+    transformer: TransformerField | None = InputField(
+        default=None,
+        description=FieldDescriptions.transformer,
+        input=Input.Connection,
+        title="Transformer",
+    )
+    qwen3_encoder: Qwen3EncoderField | None = InputField(
+        default=None,
+        title="Qwen3 Encoder",
+        description=FieldDescriptions.qwen3_encoder,
+        input=Input.Connection,
+    )
+
+    def invoke(self, context: InvocationContext) -> Flux2KleinLoRALoaderOutput:
+        lora_key = self.lora.key
+
+        if not context.models.exists(lora_key):
+            raise ValueError(f"Unknown lora: {lora_key}!")
+
+        # Warn if LoRA variant doesn't match transformer variant
+        lora_config = context.models.get_config(lora_key)
+        lora_variant = getattr(lora_config, "variant", None)
+        if lora_variant and self.transformer is not None:
+            transformer_config = context.models.get_config(self.transformer.transformer.key)
+            transformer_variant = getattr(transformer_config, "variant", None)
+            if transformer_variant and lora_variant != transformer_variant:
+                context.logger.warning(
+                    f"LoRA variant mismatch: LoRA '{lora_config.name}' is for {lora_variant.value} "
+                    f"but transformer is {transformer_variant.value}. This may cause shape errors."
+                )
+
+        # Check for existing LoRAs with the same key.
+        if self.transformer and any(lora.lora.key == lora_key for lora in self.transformer.loras):
+            raise ValueError(f'LoRA "{lora_key}" already applied to transformer.')
+        if self.qwen3_encoder and any(lora.lora.key == lora_key for lora in self.qwen3_encoder.loras):
+            raise ValueError(f'LoRA "{lora_key}" already applied to Qwen3 encoder.')
+
+        output = Flux2KleinLoRALoaderOutput()
+
+        # Attach LoRA layers to the models.
+        if self.transformer is not None:
+            output.transformer = self.transformer.model_copy(deep=True)
+            output.transformer.loras.append(
+                LoRAField(
+                    lora=self.lora,
+                    weight=self.weight,
+                )
+            )
+        if self.qwen3_encoder is not None:
+            output.qwen3_encoder = self.qwen3_encoder.model_copy(deep=True)
+            output.qwen3_encoder.loras.append(
+                LoRAField(
+                    lora=self.lora,
+                    weight=self.weight,
+                )
+            )
+
+        return output
+
+
+@invocation(
+    "flux2_klein_lora_collection_loader",
+    title="Apply LoRA Collection - Flux2 Klein",
+    tags=["lora", "model", "flux", "klein", "flux2"],
+    category="model",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class Flux2KleinLoRACollectionLoader(BaseInvocation):
+    """Applies a collection of LoRAs to a FLUX.2 Klein transformer and/or Qwen3 text encoder."""
+
+    loras: Optional[LoRAField | list[LoRAField]] = InputField(
+        default=None, description="LoRA models and weights. May be a single LoRA or collection.", title="LoRAs"
+    )
+
+    transformer: Optional[TransformerField] = InputField(
+        default=None,
+        description=FieldDescriptions.transformer,
+        input=Input.Connection,
+        title="Transformer",
+    )
+    qwen3_encoder: Qwen3EncoderField | None = InputField(
+        default=None,
+        title="Qwen3 Encoder",
+        description=FieldDescriptions.qwen3_encoder,
+        input=Input.Connection,
+    )
+
+    def invoke(self, context: InvocationContext) -> Flux2KleinLoRALoaderOutput:
+        output = Flux2KleinLoRALoaderOutput()
+        loras = self.loras if isinstance(self.loras, list) else [self.loras]
+        added_loras: list[str] = []
+
+        if self.transformer is not None:
+            output.transformer = self.transformer.model_copy(deep=True)
+
+        if self.qwen3_encoder is not None:
+            output.qwen3_encoder = self.qwen3_encoder.model_copy(deep=True)
+
+        for lora in loras:
+            if lora is None:
+                continue
+            if lora.lora.key in added_loras:
+                continue
+
+            if not context.models.exists(lora.lora.key):
+                raise Exception(f"Unknown lora: {lora.lora.key}!")
+
+            assert lora.lora.base in (BaseModelType.Flux, BaseModelType.Flux2)
+
+            # Warn if LoRA variant doesn't match transformer variant
+            lora_config = context.models.get_config(lora.lora.key)
+            lora_variant = getattr(lora_config, "variant", None)
+            if lora_variant and self.transformer is not None:
+                transformer_config = context.models.get_config(self.transformer.transformer.key)
+                transformer_variant = getattr(transformer_config, "variant", None)
+                if transformer_variant and lora_variant != transformer_variant:
+                    context.logger.warning(
+                        f"LoRA variant mismatch: LoRA '{lora_config.name}' is for {lora_variant.value} "
+                        f"but transformer is {transformer_variant.value}. This may cause shape errors."
+                    )
+
+            added_loras.append(lora.lora.key)
+
+            if self.transformer is not None and output.transformer is not None:
+                output.transformer.loras.append(lora)
+
+            if self.qwen3_encoder is not None and output.qwen3_encoder is not None:
+                output.qwen3_encoder.loras.append(lora)
+
+        return output

invokeai/backend/model_manager/configs/factory.py

@@ -40,11 +40,13 @@
 from invokeai.backend.model_manager.configs.llava_onevision import LlavaOnevision_Diffusers_Config
 from invokeai.backend.model_manager.configs.lora import (
     ControlLoRA_LyCORIS_FLUX_Config,
+    LoRA_Diffusers_Flux2_Config,
     LoRA_Diffusers_FLUX_Config,
     LoRA_Diffusers_SD1_Config,
     LoRA_Diffusers_SD2_Config,
     LoRA_Diffusers_SDXL_Config,
     LoRA_Diffusers_ZImage_Config,
+    LoRA_LyCORIS_Flux2_Config,
     LoRA_LyCORIS_FLUX_Config,
     LoRA_LyCORIS_SD1_Config,
     LoRA_LyCORIS_SD2_Config,
@@ -197,18 +199,24 @@
         Annotated[ControlNet_Diffusers_SDXL_Config, ControlNet_Diffusers_SDXL_Config.get_tag()],
         Annotated[ControlNet_Diffusers_FLUX_Config, ControlNet_Diffusers_FLUX_Config.get_tag()],
         # LoRA - LyCORIS format
+        # IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
+        # that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
         Annotated[LoRA_LyCORIS_SD1_Config, LoRA_LyCORIS_SD1_Config.get_tag()],
         Annotated[LoRA_LyCORIS_SD2_Config, LoRA_LyCORIS_SD2_Config.get_tag()],
         Annotated[LoRA_LyCORIS_SDXL_Config, LoRA_LyCORIS_SDXL_Config.get_tag()],
+        Annotated[LoRA_LyCORIS_Flux2_Config, LoRA_LyCORIS_Flux2_Config.get_tag()],
         Annotated[LoRA_LyCORIS_FLUX_Config, LoRA_LyCORIS_FLUX_Config.get_tag()],
         Annotated[LoRA_LyCORIS_ZImage_Config, LoRA_LyCORIS_ZImage_Config.get_tag()],
         # LoRA - OMI format
         Annotated[LoRA_OMI_SDXL_Config, LoRA_OMI_SDXL_Config.get_tag()],
         Annotated[LoRA_OMI_FLUX_Config, LoRA_OMI_FLUX_Config.get_tag()],
         # LoRA - diffusers format
+        # IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
+        # that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
         Annotated[LoRA_Diffusers_SD1_Config, LoRA_Diffusers_SD1_Config.get_tag()],
         Annotated[LoRA_Diffusers_SD2_Config, LoRA_Diffusers_SD2_Config.get_tag()],
         Annotated[LoRA_Diffusers_SDXL_Config, LoRA_Diffusers_SDXL_Config.get_tag()],
+        Annotated[LoRA_Diffusers_Flux2_Config, LoRA_Diffusers_Flux2_Config.get_tag()],
         Annotated[LoRA_Diffusers_FLUX_Config, LoRA_Diffusers_FLUX_Config.get_tag()],
         Annotated[LoRA_Diffusers_ZImage_Config, LoRA_Diffusers_ZImage_Config.get_tag()],
         # ControlLoRA - diffusers format