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