invoke-ai--invokeai/invokeai/app/invocations/z_image_text_encoder.py

from contextlib import ExitStack
from typing import Iterator, 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,
    Input,
    InputField,
    TensorField,
    UIComponent,
    ZImageConditioningField,
)
from invokeai.app.invocations.model import Qwen3EncoderField
from invokeai.app.invocations.primitives import ZImageConditioningOutput
from invokeai.app.services.shared.invocation_context import InvocationContext
from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
from invokeai.backend.patches.layer_patcher import LayerPatcher
from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_QWEN3_PREFIX
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
    ConditioningFieldData,
    ZImageConditioningInfo,
)
from invokeai.backend.util.devices import TorchDevice

# Z-Image max sequence length based on diffusers default
Z_IMAGE_MAX_SEQ_LEN = 512


@invocation(
    "z_image_text_encoder",
    title="Prompt - Z-Image",
    tags=["prompt", "conditioning", "z-image"],
    category="prompt",
    version="1.1.0",
    classification=Classification.Prototype,
)
class ZImageTextEncoderInvocation(BaseInvocation):
    """Encodes and preps a prompt for a Z-Image image.

    Supports regional prompting by connecting a mask input.
    """

    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,
    )
    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) -> ZImageConditioningOutput:
        prompt_embeds = self._encode_prompt(context, max_seq_len=Z_IMAGE_MAX_SEQ_LEN)
        # Move embeddings to CPU for storage to save VRAM
        prompt_embeds = prompt_embeds.detach().to("cpu")
        conditioning_data = ConditioningFieldData(conditionings=[ZImageConditioningInfo(prompt_embeds=prompt_embeds)])
        conditioning_name = context.conditioning.save(conditioning_data)
        return ZImageConditioningOutput(
            conditioning=ZImageConditioningField(conditioning_name=conditioning_name, mask=self.mask)
        )

    def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch.Tensor:
        """Encode prompt using Qwen3 text encoder.

        Based on the ZImagePipeline._encode_prompt method from diffusers.
        """
        prompt = self.prompt

        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())

            # Use the device that the text encoder is effectively executing on, and repair any required tensors left on
            # the CPU by a previous interrupted run.
            repaired_tensors = text_encoder_info.repair_required_tensors_on_device()
            device = get_effective_device(text_encoder)
            if repaired_tensors > 0:
                context.logger.warning(
                    f"Recovered {repaired_tensors} required Qwen3 tensor(s) onto {device} after a partial device mismatch."
                )

            # 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=Z_IMAGE_LORA_QWEN3_PREFIX,
                    dtype=lora_dtype,
                    cached_weights=cached_weights,
                )
            )

            context.util.signal_progress("Running Qwen3 text encoder")
            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."
                )

            # Apply chat template similar to diffusers ZImagePipeline
            # The chat template formats the prompt for the Qwen3 model
            try:
                prompt_formatted = tokenizer.apply_chat_template(
                    [{"role": "user", "content": prompt}],
                    tokenize=False,
                    add_generation_prompt=True,
                    enable_thinking=True,
                )
            except (AttributeError, TypeError) as e:
                # Fallback if tokenizer doesn't support apply_chat_template or enable_thinking
                context.logger.warning(f"Chat template failed ({e}), using raw prompt.")
                prompt_formatted = prompt

            # Tokenize the formatted prompt
            text_inputs = tokenizer(
                prompt_formatted,
                padding="max_length",
                max_length=max_seq_len,
                truncation=True,
                return_attention_mask=True,
                return_tensors="pt",
            )

            text_input_ids = text_inputs.input_ids
            attention_mask = text_inputs.attention_mask
            if not isinstance(text_input_ids, torch.Tensor):
                raise TypeError(
                    f"Expected torch.Tensor for input_ids, got {type(text_input_ids).__name__}. "
                    "Tokenizer returned unexpected type."
                )
            if not isinstance(attention_mask, torch.Tensor):
                raise TypeError(
                    f"Expected torch.Tensor for attention_mask, got {type(attention_mask).__name__}. "
                    "Tokenizer returned unexpected type."
                )

            # Check for truncation
            untruncated_ids = tokenizer(prompt_formatted, padding="longest", return_tensors="pt").input_ids
            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
                text_input_ids, untruncated_ids
            ):
                removed_text = tokenizer.batch_decode(untruncated_ids[:, max_seq_len - 1 : -1])
                context.logger.warning(
                    f"The following part of your input was truncated because `max_sequence_length` is set to "
                    f"{max_seq_len} tokens: {removed_text}"
                )

            # Get hidden states from the text encoder
            # Use the second-to-last hidden state like diffusers does
            prompt_mask = attention_mask.to(device).bool()
            outputs = text_encoder(
                text_input_ids.to(device),
                attention_mask=prompt_mask,
                output_hidden_states=True,
            )

            # 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."
                )
            if len(outputs.hidden_states) < 2:
                raise RuntimeError(
                    f"Expected at least 2 hidden states from text encoder, got {len(outputs.hidden_states)}. "
                    "This may indicate an incompatible model or configuration."
                )
            prompt_embeds = outputs.hidden_states[-2]

            # Z-Image expects a 2D tensor [seq_len, hidden_dim] with only valid tokens
            # Based on diffusers ZImagePipeline implementation:
            # embeddings_list.append(prompt_embeds[i][prompt_masks[i]])
            # Since batch_size=1, we take the first item and filter by mask
            prompt_embeds = prompt_embeds[0][prompt_mask[0]]

        if not isinstance(prompt_embeds, torch.Tensor):
            raise TypeError(
                f"Expected torch.Tensor for prompt embeddings, got {type(prompt_embeds).__name__}. "
                "Text encoder returned unexpected type."
            )
        return prompt_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