chore: import upstream snapshot with attribution

invokeai/app/invocations/anima_text_encoder.py

@@ -0,0 +1,216 @@
+"""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