"""Anima text encoder invocation. Encodes text using the dual-conditioning pipeline: 1. Qwen3 0.6B: Produces hidden states (last layer) 2. T5-XXL Tokenizer: Produces token IDs only (no T5 model needed) Both outputs are stored together in AnimaConditioningInfo and used by the LLM Adapter inside the transformer during denoising. Key differences from Z-Image text encoder: - Anima uses Qwen3 0.6B (base model, NOT instruct) — no chat template - Anima additionally tokenizes with T5-XXL tokenizer to get token IDs - Qwen3 output uses all positions (including padding) for full context """ from contextlib import ExitStack from typing import Iterator, Tuple import torch from transformers import PreTrainedModel, PreTrainedTokenizerBase from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation from invokeai.app.invocations.fields import ( AnimaConditioningField, FieldDescriptions, Input, InputField, TensorField, UIComponent, ) from invokeai.app.invocations.model import Qwen3EncoderField from invokeai.app.invocations.primitives import AnimaConditioningOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.anima.t5_tokenizer import load_bundled_t5_tokenizer from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.lora_conversions.anima_lora_constants import ANIMA_LORA_QWEN3_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( AnimaConditioningInfo, ConditioningFieldData, ) from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger logger = InvokeAILogger.get_logger(__name__) # T5-XXL max sequence length for token IDs T5_MAX_SEQ_LEN = 512 # Safety cap for Qwen3 sequence length to prevent GPU OOM on extremely long prompts. # Qwen3 0.6B supports 32K context but the LLM Adapter doesn't need that much. QWEN3_MAX_SEQ_LEN = 8192 @invocation( "anima_text_encoder", title="Prompt - Anima", tags=["prompt", "conditioning", "anima"], category="conditioning", version="1.4.0", classification=Classification.Prototype, ) class AnimaTextEncoderInvocation(BaseInvocation): """Encodes and preps a prompt for an Anima image. Uses Qwen3 0.6B for hidden state extraction and a bundled T5-XXL tokenizer for token IDs (no T5 model weights needed). Both are combined by the LLM Adapter inside the Anima transformer during denoising. """ 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: TensorField | None = InputField( default=None, description="A mask defining the region that this conditioning prompt applies to.", ) @torch.no_grad() def invoke(self, context: InvocationContext) -> AnimaConditioningOutput: qwen3_embeds, t5xxl_ids, t5xxl_weights = self._encode_prompt(context) # Move to CPU for storage qwen3_embeds = qwen3_embeds.detach().to("cpu") t5xxl_ids = t5xxl_ids.detach().to("cpu") t5xxl_weights = t5xxl_weights.detach().to("cpu") if t5xxl_weights is not None else None conditioning_data = ConditioningFieldData( conditionings=[ AnimaConditioningInfo( qwen3_embeds=qwen3_embeds, t5xxl_ids=t5xxl_ids, t5xxl_weights=t5xxl_weights, ) ] ) conditioning_name = context.conditioning.save(conditioning_data) return AnimaConditioningOutput( conditioning=AnimaConditioningField(conditioning_name=conditioning_name, mask=self.mask) ) def _encode_prompt( self, context: InvocationContext, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """Encode prompt using Qwen3 0.6B and T5-XXL tokenizer. Returns: Tuple of (qwen3_embeds, t5xxl_ids, t5xxl_weights). - qwen3_embeds: Shape (max_seq_len, 1024) — includes all positions (including padding) to preserve full sequence context for the LLM Adapter. - t5xxl_ids: Shape (seq_len,) — T5-XXL token IDs (unpadded). - t5xxl_weights: None (uniform weights for now). """ prompt = self.prompt # --- Step 1: Encode with Qwen3 0.6B --- 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: (_, text_encoder) = exit_stack.enter_context(text_encoder_info.model_on_device()) (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device()) device = text_encoder.device # Apply LoRA models to the text encoder lora_dtype = TorchDevice.choose_anima_inference_dtype(device) exit_stack.enter_context( LayerPatcher.apply_smart_model_patches( model=text_encoder, patches=self._lora_iterator(context), prefix=ANIMA_LORA_QWEN3_PREFIX, dtype=lora_dtype, ) ) if not isinstance(text_encoder, PreTrainedModel): raise TypeError(f"Expected PreTrainedModel for text encoder, got {type(text_encoder).__name__}.") if not isinstance(tokenizer, PreTrainedTokenizerBase): raise TypeError(f"Expected PreTrainedTokenizerBase for tokenizer, got {type(tokenizer).__name__}.") context.util.signal_progress("Running Qwen3 0.6B text encoder") # Anima uses base Qwen3 (not instruct) — tokenize directly, no chat template. # A safety cap is applied to prevent GPU OOM on extremely long prompts. text_inputs = tokenizer( prompt, padding=False, truncation=True, max_length=QWEN3_MAX_SEQ_LEN, 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) or not isinstance(attention_mask, torch.Tensor): raise TypeError("Tokenizer returned unexpected types.") if text_input_ids.shape[-1] == QWEN3_MAX_SEQ_LEN: logger.warning( f"Prompt was truncated to {QWEN3_MAX_SEQ_LEN} tokens. " "Consider shortening the prompt for best results." ) # Ensure at least 1 token (empty prompts produce 0 tokens with padding=False) if text_input_ids.shape[-1] == 0: pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id text_input_ids = torch.tensor([[pad_id]]) attention_mask = torch.tensor([[1]]) # Get last hidden state from Qwen3 (final layer output) prompt_mask = attention_mask.to(device).bool() outputs = text_encoder( text_input_ids.to(device), attention_mask=prompt_mask, output_hidden_states=True, ) if not hasattr(outputs, "hidden_states") or outputs.hidden_states is None: raise RuntimeError("Text encoder did not return hidden_states.") if len(outputs.hidden_states) < 1: raise RuntimeError(f"Expected at least 1 hidden state, got {len(outputs.hidden_states)}.") # Use last hidden state — only real tokens, no padding qwen3_embeds = outputs.hidden_states[-1][0] # Shape: (seq_len, 1024) # --- Step 2: Tokenize with bundled T5-XXL tokenizer (IDs only, no model) --- context.util.signal_progress("Tokenizing with T5-XXL") t5_tokenizer = load_bundled_t5_tokenizer() t5_tokens = t5_tokenizer( prompt, padding=False, truncation=True, max_length=T5_MAX_SEQ_LEN, return_tensors="pt", ) t5xxl_ids = t5_tokens.input_ids[0] # Shape: (seq_len,) return qwen3_embeds, t5xxl_ids, None 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