InvokeAI 6.10.0rc1__py3-none-any.whl → 6.11.0__py3-none-any.whl

invokeai/app/invocations/flux2_klein_model_loader.py

@@ -0,0 +1,222 @@
+"""Flux2 Klein Model Loader Invocation.
+
+Loads a Flux2 Klein model with its Qwen3 text encoder and VAE.
+Unlike standard FLUX which uses CLIP+T5, Klein uses only Qwen3.
+"""
+
+from typing import Literal, 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 (
+    ModelIdentifierField,
+    Qwen3EncoderField,
+    TransformerField,
+    VAEField,
+)
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    Flux2VariantType,
+    ModelFormat,
+    ModelType,
+    Qwen3VariantType,
+    SubModelType,
+)
+
+
+@invocation_output("flux2_klein_model_loader_output")
+class Flux2KleinModelLoaderOutput(BaseInvocationOutput):
+    """Flux2 Klein model loader output."""
+
+    transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer")
+    qwen3_encoder: Qwen3EncoderField = OutputField(description=FieldDescriptions.qwen3_encoder, title="Qwen3 Encoder")
+    vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE")
+    max_seq_len: Literal[256, 512] = OutputField(
+        description="The max sequence length for the Qwen3 encoder.",
+        title="Max Seq Length",
+    )
+
+
+@invocation(
+    "flux2_klein_model_loader",
+    title="Main Model - Flux2 Klein",
+    tags=["model", "flux", "klein", "qwen3"],
+    category="model",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class Flux2KleinModelLoaderInvocation(BaseInvocation):
+    """Loads a Flux2 Klein model, outputting its submodels.
+
+    Flux2 Klein uses Qwen3 as the text encoder instead of CLIP+T5.
+    It uses a 32-channel VAE (AutoencoderKLFlux2) instead of the 16-channel FLUX.1 VAE.
+
+    When using a Diffusers format model, both VAE and Qwen3 encoder are extracted
+    automatically from the main model. You can override with standalone models:
+    - Transformer: Always from Flux2 Klein main model
+    - VAE: From main model (Diffusers) or standalone VAE
+    - Qwen3 Encoder: From main model (Diffusers) or standalone Qwen3 model
+    """
+
+    model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.flux_model,
+        input=Input.Direct,
+        ui_model_base=BaseModelType.Flux2,
+        ui_model_type=ModelType.Main,
+        title="Transformer",
+    )
+
+    vae_model: Optional[ModelIdentifierField] = InputField(
+        default=None,
+        description="Standalone VAE model. Flux2 Klein uses the same VAE as FLUX (16-channel). "
+        "If not provided, VAE will be loaded from the Qwen3 Source model.",
+        input=Input.Direct,
+        ui_model_base=[BaseModelType.Flux, BaseModelType.Flux2],
+        ui_model_type=ModelType.VAE,
+        title="VAE",
+    )
+
+    qwen3_encoder_model: Optional[ModelIdentifierField] = InputField(
+        default=None,
+        description="Standalone Qwen3 Encoder model. "
+        "If not provided, encoder will be loaded from the Qwen3 Source model.",
+        input=Input.Direct,
+        ui_model_type=ModelType.Qwen3Encoder,
+        title="Qwen3 Encoder",
+    )
+
+    qwen3_source_model: Optional[ModelIdentifierField] = InputField(
+        default=None,
+        description="Diffusers Flux2 Klein model to extract VAE and/or Qwen3 encoder from. "
+        "Use this if you don't have separate VAE/Qwen3 models. "
+        "Ignored if both VAE and Qwen3 Encoder are provided separately.",
+        input=Input.Direct,
+        ui_model_base=BaseModelType.Flux2,
+        ui_model_type=ModelType.Main,
+        ui_model_format=ModelFormat.Diffusers,
+        title="Qwen3 Source (Diffusers)",
+    )
+
+    max_seq_len: Literal[256, 512] = InputField(
+        default=512,
+        description="Max sequence length for the Qwen3 encoder.",
+        title="Max Seq Length",
+    )
+
+    def invoke(self, context: InvocationContext) -> Flux2KleinModelLoaderOutput:
+        # Transformer always comes from the main model
+        transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
+
+        # Check if main model is Diffusers format (can extract VAE directly)
+        main_config = context.models.get_config(self.model)
+        main_is_diffusers = main_config.format == ModelFormat.Diffusers
+
+        # Determine VAE source
+        # IMPORTANT: FLUX.2 Klein uses a 32-channel VAE (AutoencoderKLFlux2), not the 16-channel FLUX.1 VAE.
+        # The VAE should come from the FLUX.2 Klein Diffusers model, not a separate FLUX VAE.
+        if self.vae_model is not None:
+            # Use standalone VAE (user explicitly selected one)
+            vae = self.vae_model.model_copy(update={"submodel_type": SubModelType.VAE})
+        elif main_is_diffusers:
+            # Extract VAE from main model (recommended for FLUX.2)
+            vae = self.model.model_copy(update={"submodel_type": SubModelType.VAE})
+        elif self.qwen3_source_model is not None:
+            # Extract from Qwen3 source Diffusers model
+            self._validate_diffusers_format(context, self.qwen3_source_model, "Qwen3 Source")
+            vae = self.qwen3_source_model.model_copy(update={"submodel_type": SubModelType.VAE})
+        else:
+            raise ValueError(
+                "No VAE source provided. Standalone safetensors/GGUF models require a separate VAE. "
+                "Options:\n"
+                "  1. Set 'VAE' to a standalone FLUX VAE model\n"
+                "  2. Set 'Qwen3 Source' to a Diffusers Flux2 Klein model to extract the VAE from"
+            )
+
+        # Determine Qwen3 Encoder source
+        if self.qwen3_encoder_model is not None:
+            # Use standalone Qwen3 Encoder - validate it matches the FLUX.2 Klein variant
+            self._validate_qwen3_encoder_variant(context, main_config)
+            qwen3_tokenizer = self.qwen3_encoder_model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+            qwen3_encoder = self.qwen3_encoder_model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+        elif main_is_diffusers:
+            # Extract from main model (recommended for FLUX.2 Klein)
+            qwen3_tokenizer = self.model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+            qwen3_encoder = self.model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+        elif self.qwen3_source_model is not None:
+            # Extract from separate Diffusers model
+            self._validate_diffusers_format(context, self.qwen3_source_model, "Qwen3 Source")
+            qwen3_tokenizer = self.qwen3_source_model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+            qwen3_encoder = self.qwen3_source_model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+        else:
+            raise ValueError(
+                "No Qwen3 Encoder source provided. Standalone safetensors/GGUF models require a separate text encoder. "
+                "Options:\n"
+                "  1. Set 'Qwen3 Encoder' to a standalone Qwen3 text encoder model "
+                "(Klein 4B needs Qwen3 4B, Klein 9B needs Qwen3 8B)\n"
+                "  2. Set 'Qwen3 Source' to a Diffusers Flux2 Klein model to extract the encoder from"
+            )
+
+        return Flux2KleinModelLoaderOutput(
+            transformer=TransformerField(transformer=transformer, loras=[]),
+            qwen3_encoder=Qwen3EncoderField(tokenizer=qwen3_tokenizer, text_encoder=qwen3_encoder),
+            vae=VAEField(vae=vae),
+            max_seq_len=self.max_seq_len,
+        )
+
+    def _validate_diffusers_format(
+        self, context: InvocationContext, model: ModelIdentifierField, model_name: str
+    ) -> None:
+        """Validate that a model is in Diffusers format."""
+        config = context.models.get_config(model)
+        if config.format != ModelFormat.Diffusers:
+            raise ValueError(
+                f"The {model_name} model must be a Diffusers format model. "
+                f"The selected model '{config.name}' is in {config.format.value} format."
+            )
+
+    def _validate_qwen3_encoder_variant(self, context: InvocationContext, main_config) -> None:
+        """Validate that the standalone Qwen3 encoder variant matches the FLUX.2 Klein variant.
+
+        - FLUX.2 Klein 4B requires Qwen3 4B encoder
+        - FLUX.2 Klein 9B requires Qwen3 8B encoder
+        """
+        if self.qwen3_encoder_model is None:
+            return
+
+        # Get the Qwen3 encoder config
+        qwen3_config = context.models.get_config(self.qwen3_encoder_model)
+
+        # Check if the config has a variant field
+        if not hasattr(qwen3_config, "variant"):
+            # Can't validate, skip
+            return
+
+        qwen3_variant = qwen3_config.variant
+
+        # Get the FLUX.2 Klein variant from the main model config
+        if not hasattr(main_config, "variant"):
+            return
+
+        flux2_variant = main_config.variant
+
+        # Validate the variants match
+        # Klein4B requires Qwen3_4B, Klein9B/Klein9BBase requires Qwen3_8B
+        expected_qwen3_variant = None
+        if flux2_variant == Flux2VariantType.Klein4B:
+            expected_qwen3_variant = Qwen3VariantType.Qwen3_4B
+        elif flux2_variant in (Flux2VariantType.Klein9B, Flux2VariantType.Klein9BBase):
+            expected_qwen3_variant = Qwen3VariantType.Qwen3_8B
+
+        if expected_qwen3_variant is not None and qwen3_variant != expected_qwen3_variant:
+            raise ValueError(
+                f"Qwen3 encoder variant mismatch: FLUX.2 Klein {flux2_variant.value} requires "
+                f"{expected_qwen3_variant.value} encoder, but {qwen3_variant.value} was selected. "
+                f"Please select a matching Qwen3 encoder or use a Diffusers format model which includes the correct encoder."
+            )

invokeai/app/invocations/flux2_klein_text_encoder.py

@@ -0,0 +1,222 @@
+"""Flux2 Klein Text Encoder Invocation.
+
+Flux2 Klein uses Qwen3 as the text encoder instead of CLIP+T5.
+The key difference is that it extracts hidden states from layers (9, 18, 27)
+and stacks them together for richer text representations.
+
+This implementation matches the diffusers Flux2KleinPipeline exactly.
+"""
+
+from contextlib import ExitStack
+from typing import Iterator, Literal, Optional, Tuple
+
+import torch
+from transformers import PreTrainedModel, PreTrainedTokenizerBase
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    FluxConditioningField,
+    Input,
+    InputField,
+    TensorField,
+    UIComponent,
+)
+from invokeai.app.invocations.model import Qwen3EncoderField
+from invokeai.app.invocations.primitives import FluxConditioningOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.patches.layer_patcher import LayerPatcher
+from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_T5_PREFIX
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, FLUXConditioningInfo
+from invokeai.backend.util.devices import TorchDevice
+
+# FLUX.2 Klein extracts hidden states from these specific layers
+# Matching diffusers Flux2KleinPipeline: (9, 18, 27)
+# hidden_states[0] is embedding layer, so layer N is at index N
+KLEIN_EXTRACTION_LAYERS = (9, 18, 27)
+
+# Default max sequence length for Klein models
+KLEIN_MAX_SEQ_LEN = 512
+
+
+@invocation(
+    "flux2_klein_text_encoder",
+    title="Prompt - Flux2 Klein",
+    tags=["prompt", "conditioning", "flux", "klein", "qwen3"],
+    category="conditioning",
+    version="1.1.0",
+    classification=Classification.Prototype,
+)
+class Flux2KleinTextEncoderInvocation(BaseInvocation):
+    """Encodes and preps a prompt for Flux2 Klein image generation.
+
+    Flux2 Klein uses Qwen3 as the text encoder, extracting hidden states from
+    layers (9, 18, 27) and stacking them for richer text representations.
+    This matches the diffusers Flux2KleinPipeline implementation exactly.
+    """
+
+    prompt: str = InputField(description="Text prompt to encode.", ui_component=UIComponent.Textarea)
+    qwen3_encoder: Qwen3EncoderField = InputField(
+        title="Qwen3 Encoder",
+        description=FieldDescriptions.qwen3_encoder,
+        input=Input.Connection,
+    )
+    max_seq_len: Literal[256, 512] = InputField(
+        default=512,
+        description="Max sequence length for the Qwen3 encoder.",
+    )
+    mask: Optional[TensorField] = InputField(
+        default=None,
+        description="A mask defining the region that this conditioning prompt applies to.",
+    )
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> FluxConditioningOutput:
+        qwen3_embeds, pooled_embeds = self._encode_prompt(context)
+
+        # Use FLUXConditioningInfo for compatibility with existing Flux denoiser
+        # t5_embeds -> qwen3 stacked embeddings
+        # clip_embeds -> pooled qwen3 embedding
+        conditioning_data = ConditioningFieldData(
+            conditionings=[FLUXConditioningInfo(clip_embeds=pooled_embeds, t5_embeds=qwen3_embeds)]
+        )
+
+        conditioning_name = context.conditioning.save(conditioning_data)
+        return FluxConditioningOutput(
+            conditioning=FluxConditioningField(conditioning_name=conditioning_name, mask=self.mask)
+        )
+
+    def _encode_prompt(self, context: InvocationContext) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Encode prompt using Qwen3 text encoder with Klein-style layer extraction.
+
+        This matches the diffusers Flux2KleinPipeline._get_qwen3_prompt_embeds() exactly.
+
+        Returns:
+            Tuple of (stacked_embeddings, pooled_embedding):
+            - stacked_embeddings: Hidden states from layers (9, 18, 27) stacked together.
+              Shape: (1, seq_len, hidden_size * 3)
+            - pooled_embedding: Pooled representation for global conditioning.
+              Shape: (1, hidden_size)
+        """
+        prompt = self.prompt
+        device = TorchDevice.choose_torch_device()
+
+        text_encoder_info = context.models.load(self.qwen3_encoder.text_encoder)
+        tokenizer_info = context.models.load(self.qwen3_encoder.tokenizer)
+
+        with ExitStack() as exit_stack:
+            (cached_weights, text_encoder) = exit_stack.enter_context(text_encoder_info.model_on_device())
+            (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device())
+
+            # Apply LoRA models to the text encoder
+            lora_dtype = TorchDevice.choose_bfloat16_safe_dtype(device)
+            exit_stack.enter_context(
+                LayerPatcher.apply_smart_model_patches(
+                    model=text_encoder,
+                    patches=self._lora_iterator(context),
+                    prefix=FLUX_LORA_T5_PREFIX,  # Reuse T5 prefix for Qwen3 LoRAs
+                    dtype=lora_dtype,
+                    cached_weights=cached_weights,
+                )
+            )
+
+            context.util.signal_progress("Running Qwen3 text encoder (Klein)")
+
+            if not isinstance(text_encoder, PreTrainedModel):
+                raise TypeError(
+                    f"Expected PreTrainedModel for text encoder, got {type(text_encoder).__name__}. "
+                    "The Qwen3 encoder model may be corrupted or incompatible."
+                )
+            if not isinstance(tokenizer, PreTrainedTokenizerBase):
+                raise TypeError(
+                    f"Expected PreTrainedTokenizerBase for tokenizer, got {type(tokenizer).__name__}. "
+                    "The Qwen3 tokenizer may be corrupted or incompatible."
+                )
+
+            # Format messages exactly like diffusers Flux2KleinPipeline:
+            # - Only user message, NO system message
+            # - add_generation_prompt=True (adds assistant prefix)
+            # - enable_thinking=False
+            messages = [{"role": "user", "content": prompt}]
+
+            # Step 1: Apply chat template to get formatted text (tokenize=False)
+            text: str = tokenizer.apply_chat_template(  # type: ignore[assignment]
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,  # Adds assistant prefix like diffusers
+                enable_thinking=False,  # Disable thinking mode
+            )
+
+            # Step 2: Tokenize the formatted text
+            inputs = tokenizer(
+                text,
+                return_tensors="pt",
+                padding="max_length",
+                truncation=True,
+                max_length=self.max_seq_len,
+            )
+
+            input_ids = inputs["input_ids"]
+            attention_mask = inputs["attention_mask"]
+
+            # Move to device
+            input_ids = input_ids.to(device)
+            attention_mask = attention_mask.to(device)
+
+            # Forward pass through the model - matching diffusers exactly
+            outputs = text_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+                use_cache=False,
+            )
+
+            # Validate hidden_states output
+            if not hasattr(outputs, "hidden_states") or outputs.hidden_states is None:
+                raise RuntimeError(
+                    "Text encoder did not return hidden_states. "
+                    "Ensure output_hidden_states=True is supported by this model."
+                )
+
+            num_hidden_layers = len(outputs.hidden_states)
+
+            # Extract and stack hidden states - EXACTLY like diffusers:
+            # out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1)
+            # prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
+            hidden_states_list = []
+            for layer_idx in KLEIN_EXTRACTION_LAYERS:
+                if layer_idx >= num_hidden_layers:
+                    layer_idx = num_hidden_layers - 1
+                hidden_states_list.append(outputs.hidden_states[layer_idx])
+
+            # Stack along dim=1, then permute and reshape - exactly like diffusers
+            out = torch.stack(hidden_states_list, dim=1)
+            out = out.to(dtype=text_encoder.dtype, device=device)
+
+            batch_size, num_channels, seq_len, hidden_dim = out.shape
+            prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
+
+            # Create pooled embedding for global conditioning
+            # Use mean pooling over the sequence (excluding padding)
+            # This serves a similar role to CLIP's pooled output in standard FLUX
+            last_hidden_state = outputs.hidden_states[-1]  # Use last layer for pooling
+            # Expand mask to match hidden state dimensions
+            expanded_mask = attention_mask.unsqueeze(-1).expand_as(last_hidden_state).float()
+            sum_embeds = (last_hidden_state * expanded_mask).sum(dim=1)
+            num_tokens = expanded_mask.sum(dim=1).clamp(min=1)
+            pooled_embeds = sum_embeds / num_tokens
+
+        return prompt_embeds, pooled_embeds
+
+    def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]:
+        """Iterate over LoRA models to apply to the Qwen3 text encoder."""
+        for lora in self.qwen3_encoder.loras:
+            lora_info = context.models.load(lora.lora)
+            if not isinstance(lora_info.model, ModelPatchRaw):
+                raise TypeError(
+                    f"Expected ModelPatchRaw for LoRA '{lora.lora.key}', got {type(lora_info.model).__name__}. "
+                    "The LoRA model may be corrupted or incompatible."
+                )
+            yield (lora_info.model, lora.weight)
+            del lora_info

invokeai/app/invocations/flux2_vae_decode.py

@@ -0,0 +1,106 @@
+"""Flux2 Klein VAE Decode Invocation.
+
+Decodes latents to images using the FLUX.2 32-channel VAE (AutoencoderKLFlux2).
+"""
+
+import torch
+from einops import rearrange
+from PIL import Image
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    Input,
+    InputField,
+    LatentsField,
+    WithBoard,
+    WithMetadata,
+)
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.invocations.primitives import ImageOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.load.load_base import LoadedModel
+from invokeai.backend.util.devices import TorchDevice
+
+
+@invocation(
+    "flux2_vae_decode",
+    title="Latents to Image - FLUX2",
+    tags=["latents", "image", "vae", "l2i", "flux2", "klein"],
+    category="latents",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class Flux2VaeDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Generates an image from latents using FLUX.2 Klein's 32-channel VAE."""
+
+    latents: LatentsField = InputField(
+        description=FieldDescriptions.latents,
+        input=Input.Connection,
+    )
+    vae: VAEField = InputField(
+        description=FieldDescriptions.vae,
+        input=Input.Connection,
+    )
+
+    def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Image:
+        """Decode latents to image using FLUX.2 VAE.
+
+        Input latents should already be in the correct space after BN denormalization
+        was applied in the denoiser. The VAE expects (B, 32, H, W) format.
+        """
+        with vae_info.model_on_device() as (_, vae):
+            vae_dtype = next(iter(vae.parameters())).dtype
+            device = TorchDevice.choose_torch_device()
+            latents = latents.to(device=device, dtype=vae_dtype)
+
+            # Decode using diffusers API
+            decoded = vae.decode(latents, return_dict=False)[0]
+
+        # Debug: Log decoded output statistics
+        print(
+            f"[FLUX.2 VAE] Decoded output: shape={decoded.shape}, "
+            f"min={decoded.min().item():.4f}, max={decoded.max().item():.4f}, "
+            f"mean={decoded.mean().item():.4f}"
+        )
+        # Check per-channel statistics to diagnose color issues
+        for c in range(min(3, decoded.shape[1])):
+            ch = decoded[0, c]
+            print(
+                f"[FLUX.2 VAE] Channel {c}: min={ch.min().item():.4f}, "
+                f"max={ch.max().item():.4f}, mean={ch.mean().item():.4f}"
+            )
+
+        # Convert from [-1, 1] to [0, 1] then to [0, 255] PIL image
+        img = (decoded / 2 + 0.5).clamp(0, 1)
+        img = rearrange(img[0], "c h w -> h w c")
+        img_np = (img * 255).byte().cpu().numpy()
+        # Explicitly create RGB image (not grayscale)
+        img_pil = Image.fromarray(img_np, mode="RGB")
+        return img_pil
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        latents = context.tensors.load(self.latents.latents_name)
+
+        # Log latent statistics for debugging black image issues
+        context.logger.debug(
+            f"FLUX.2 VAE decode input: shape={latents.shape}, "
+            f"min={latents.min().item():.4f}, max={latents.max().item():.4f}, "
+            f"mean={latents.mean().item():.4f}"
+        )
+
+        # Warn if input latents are all zeros or very small (would cause black images)
+        if latents.abs().max() < 1e-6:
+            context.logger.warning(
+                "FLUX.2 VAE decode received near-zero latents! This will cause black images. "
+                "The latent cache may be corrupted - try clearing the cache."
+            )
+
+        vae_info = context.models.load(self.vae.vae)
+        context.util.signal_progress("Running VAE")
+        image = self._vae_decode(vae_info=vae_info, latents=latents)
+
+        TorchDevice.empty_cache()
+        image_dto = context.images.save(image=image)
+        return ImageOutput.build(image_dto)

invokeai/app/invocations/flux2_vae_encode.py

@@ -0,0 +1,88 @@
+"""Flux2 Klein VAE Encode Invocation.
+
+Encodes images to latents using the FLUX.2 32-channel VAE (AutoencoderKLFlux2).
+"""
+
+import einops
+import torch
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    ImageField,
+    Input,
+    InputField,
+)
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.invocations.primitives import LatentsOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.load.load_base import LoadedModel
+from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
+from invokeai.backend.util.devices import TorchDevice
+
+
+@invocation(
+    "flux2_vae_encode",
+    title="Image to Latents - FLUX2",
+    tags=["latents", "image", "vae", "i2l", "flux2", "klein"],
+    category="latents",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class Flux2VaeEncodeInvocation(BaseInvocation):
+    """Encodes an image into latents using FLUX.2 Klein's 32-channel VAE."""
+
+    image: ImageField = InputField(
+        description="The image to encode.",
+    )
+    vae: VAEField = InputField(
+        description=FieldDescriptions.vae,
+        input=Input.Connection,
+    )
+
+    def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor:
+        """Encode image to latents using FLUX.2 VAE.
+
+        The VAE encodes to 32-channel latent space.
+        Output latents shape: (B, 32, H/8, W/8).
+        """
+        with vae_info.model_on_device() as (_, vae):
+            vae_dtype = next(iter(vae.parameters())).dtype
+            device = TorchDevice.choose_torch_device()
+            image_tensor = image_tensor.to(device=device, dtype=vae_dtype)
+
+            # Encode using diffusers API
+            # The VAE.encode() returns a DiagonalGaussianDistribution-like object
+            latent_dist = vae.encode(image_tensor, return_dict=False)[0]
+
+            # Sample from the distribution (or use mode for deterministic output)
+            # Using mode() for deterministic encoding
+            if hasattr(latent_dist, "mode"):
+                latents = latent_dist.mode()
+            elif hasattr(latent_dist, "sample"):
+                # Fall back to sampling if mode is not available
+                generator = torch.Generator(device=device).manual_seed(0)
+                latents = latent_dist.sample(generator=generator)
+            else:
+                # Direct tensor output (some VAE implementations)
+                latents = latent_dist
+
+        return latents
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> LatentsOutput:
+        image = context.images.get_pil(self.image.image_name)
+
+        vae_info = context.models.load(self.vae.vae)
+
+        # Convert image to tensor (HWC -> CHW, normalize to [-1, 1])
+        image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
+        if image_tensor.dim() == 3:
+            image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
+
+        context.util.signal_progress("Running VAE Encode")
+        latents = self._vae_encode(vae_info=vae_info, image_tensor=image_tensor)
+
+        latents = latents.to("cpu")
+        name = context.tensors.save(tensor=latents)
+        return LatentsOutput.build(latents_name=name, latents=latents, seed=None)