feat(anima): add prompt weighting support

invokeai/app/invocations/anima_text_encoder.py

@@ -31,6 +31,7 @@
 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.prompt_weighting import parse_prompt_attention, tokenize_t5_with_weights
 from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec
 from invokeai.backend.patches.lora_conversions.anima_lora_constants import ANIMA_LORA_QWEN3_PREFIX
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
@@ -106,17 +107,17 @@ def invoke(self, context: InvocationContext) -> AnimaConditioningOutput:
     def _encode_prompt(
         self,
         context: InvocationContext,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """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).
+            - t5xxl_weights: Shape (seq_len,) — per-token prompt weights.
         """
-        prompt = self.prompt
+        prompt, weighted_ranges = parse_prompt_attention(self.prompt)
 
         # --- Step 1: Encode with Qwen3 0.6B ---
         text_encoder_info = context.models.load(self.qwen3_encoder.text_encoder)
@@ -196,16 +197,11 @@ def _encode_prompt(
         # --- 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, t5xxl_weights = tokenize_t5_with_weights(
+            t5_tokenizer, prompt, weighted_ranges, max_length=T5_MAX_SEQ_LEN
         )
-        t5xxl_ids = t5_tokens.input_ids[0]  # Shape: (seq_len,)
 
-        return qwen3_embeds, t5xxl_ids, None
+        return qwen3_embeds, t5xxl_ids, t5xxl_weights
 
     def _lora_iterator(self, context: InvocationContext) -> Iterator[PatchSpec]:
         """Iterate over LoRA models to apply to the Qwen3 text encoder."""

invokeai/backend/anima/prompt_weighting.py

@@ -0,0 +1,164 @@
+"""Prompt-weighting helpers for Anima text conditioning."""
+
+import math
+
+import torch
+from transformers import PreTrainedTokenizerBase
+
+WeightedTextRange = tuple[int, int, float]
+
+
+def _find_closing_parenthesis(prompt: str, opening_index: int) -> int | None:
+    depth = 1
+    index = opening_index + 1
+    while index < len(prompt):
+        if prompt[index] == "\\" and index + 1 < len(prompt) and prompt[index + 1] in "()\\":
+            index += 2
+            continue
+        if prompt[index] == "(":
+            depth += 1
+        elif prompt[index] == ")":
+            depth -= 1
+            if depth == 0:
+                return index
+        index += 1
+    return None
+
+
+def _split_explicit_weight(group: str) -> tuple[str, float] | None:
+    """Split a trailing top-level ``:weight`` from a parenthesized group."""
+    depth = 0
+    last_colon: int | None = None
+    index = 0
+    while index < len(group):
+        if group[index] == "\\" and index + 1 < len(group) and group[index + 1] in "()\\":
+            index += 2
+            continue
+        if group[index] == "(":
+            depth += 1
+        elif group[index] == ")":
+            depth -= 1
+        elif group[index] == ":" and depth == 0:
+            last_colon = index
+        index += 1
+
+    if last_colon is None or last_colon == 0:
+        return None
+
+    try:
+        weight = float(group[last_colon + 1 :])
+    except ValueError:
+        return None
+    if not math.isfinite(weight):
+        return None
+    return group[:last_colon], weight
+
+
+def parse_prompt_attention(prompt: str) -> tuple[str, list[WeightedTextRange]]:
+    """Remove prompt-weighting markup and return weighted character ranges.
+
+    Parentheses increase weight by 10%, while a trailing ``:number`` sets an
+    explicit weight for that group. Escaped parentheses remain literal text.
+    """
+    cleaned_parts: list[str] = []
+    weighted_ranges: list[WeightedTextRange] = []
+    cleaned_length = 0
+
+    def append_text(text: str, weight: float) -> None:
+        nonlocal cleaned_length
+        if not text:
+            return
+        start = cleaned_length
+        cleaned_parts.append(text)
+        cleaned_length += len(text)
+        if weighted_ranges and weighted_ranges[-1][1] == start and weighted_ranges[-1][2] == weight:
+            previous_start, _, _ = weighted_ranges[-1]
+            weighted_ranges[-1] = (previous_start, cleaned_length, weight)
+        else:
+            weighted_ranges.append((start, cleaned_length, weight))
+
+    def parse(text: str, weight: float) -> None:
+        index = 0
+        plain_text: list[str] = []
+
+        def flush_plain_text() -> None:
+            if plain_text:
+                append_text("".join(plain_text), weight)
+                plain_text.clear()
+
+        while index < len(text):
+            char = text[index]
+            if char == "\\" and index + 1 < len(text) and text[index + 1] in "()\\":
+                plain_text.append(text[index + 1])
+                index += 2
+                continue
+            if char != "(":
+                plain_text.append(char)
+                index += 1
+                continue
+
+            closing_index = _find_closing_parenthesis(text, index)
+            if closing_index is None:
+                plain_text.append(char)
+                index += 1
+                continue
+
+            flush_plain_text()
+            group = text[index + 1 : closing_index]
+            explicit_weight = _split_explicit_weight(group)
+            if explicit_weight is None:
+                parse(group, weight * 1.1)
+            else:
+                group_text, group_weight = explicit_weight
+                parse(group_text, group_weight)
+            index = closing_index + 1
+
+        flush_plain_text()
+
+    parse(prompt, 1.0)
+    return "".join(cleaned_parts), weighted_ranges
+
+
+def tokenize_t5_with_weights(
+    tokenizer: PreTrainedTokenizerBase,
+    prompt: str,
+    weighted_ranges: list[WeightedTextRange],
+    max_length: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Tokenize cleaned text once, then map character weights to T5 tokens."""
+    tokens = tokenizer(
+        prompt,
+        padding=False,
+        truncation=True,
+        max_length=max_length,
+        return_offsets_mapping=True,
+        return_special_tokens_mask=True,
+        return_tensors="pt",
+    )
+    input_ids = tokens.input_ids[0]
+    offsets = tokens.offset_mapping[0].tolist()
+    special_tokens_mask = tokens.special_tokens_mask[0].tolist()
+
+    token_weights: list[float] = []
+    range_index = 0
+    for (token_start, token_end), is_special in zip(offsets, special_tokens_mask, strict=True):
+        if is_special or token_end <= token_start:
+            token_weights.append(1.0)
+            continue
+
+        while range_index < len(weighted_ranges) and weighted_ranges[range_index][1] <= token_start:
+            range_index += 1
+
+        weighted_length = 0.0
+        covered_length = 0
+        current_range = range_index
+        while current_range < len(weighted_ranges) and weighted_ranges[current_range][0] < token_end:
+            range_start, range_end, weight = weighted_ranges[current_range]
+            overlap = max(0, min(token_end, range_end) - max(token_start, range_start))
+            weighted_length += overlap * weight
+            covered_length += overlap
+            current_range += 1
+
+        token_weights.append(weighted_length / covered_length if covered_length else 1.0)
+
+    return input_ids, torch.tensor(token_weights, dtype=torch.float32)

tests/app/invocations/test_anima_text_encoder.py

@@ -5,6 +5,8 @@
 import torch
 
 from invokeai.app.invocations.anima_text_encoder import AnimaTextEncoderInvocation
+from invokeai.backend.anima.prompt_weighting import parse_prompt_attention, tokenize_t5_with_weights
+from invokeai.backend.t5.t5_tokenizer import load_bundled_t5_tokenizer
 
 
 class FakeQwen3Encoder(torch.nn.Module):
@@ -35,8 +37,12 @@ class FakeQwen3Tokenizer:
     pad_token_id = 0
     eos_token_id = 0
 
+    def __init__(self):
+        self.prompt: str | None = None
+
     def __call__(self, prompt, **kwargs):
-        del prompt, kwargs
+        self.prompt = prompt
+        del kwargs
         return SimpleNamespace(
             input_ids=torch.tensor([[1, 2, 3]], dtype=torch.long),
             attention_mask=torch.tensor([[1, 1, 1]], dtype=torch.long),
@@ -46,7 +52,11 @@ def __call__(self, prompt, **kwargs):
 class FakeT5Tokenizer:
     def __call__(self, prompt, **kwargs):
         del prompt, kwargs
-        return SimpleNamespace(input_ids=torch.tensor([[1, 2, 3]], dtype=torch.long))
+        return SimpleNamespace(
+            input_ids=torch.tensor([[1, 2, 3]], dtype=torch.long),
+            offset_mapping=torch.tensor([[[0, 1], [1, 2], [0, 0]]], dtype=torch.long),
+            special_tokens_mask=torch.tensor([[0, 0, 1]], dtype=torch.long),
+        )
 
 
 class FakeLoadedModel:
@@ -63,7 +73,9 @@ def model_on_device(self):
         yield (None, self._model)
 
 
-def _run_encode(monkeypatch, compute_device: torch.device) -> FakeQwen3Encoder:
+def _run_encode(
+    monkeypatch, compute_device: torch.device, prompt: str = "test prompt"
+) -> tuple[FakeQwen3Encoder, FakeQwen3Tokenizer]:
     module_path = "invokeai.app.invocations.anima_text_encoder"
     text_encoder = FakeQwen3Encoder()
     tokenizer = FakeQwen3Tokenizer()
@@ -83,25 +95,86 @@ def _run_encode(monkeypatch, compute_device: torch.device) -> FakeQwen3Encoder:
     monkeypatch.setattr(f"{module_path}.load_bundled_t5_tokenizer", lambda: FakeT5Tokenizer())
 
     invocation = AnimaTextEncoderInvocation.model_construct(
-        prompt="test prompt",
+        prompt=prompt,
         qwen3_encoder=SimpleNamespace(text_encoder=SimpleNamespace(), tokenizer=SimpleNamespace(), loras=[]),
         mask=None,
     )
 
     invocation._encode_prompt(mock_context)
-    return text_encoder
+    return text_encoder, tokenizer
 
 
 def test_anima_qwen3_encode_uses_compute_device(monkeypatch):
     # Regression test for #9373: the encoder's weights are offloaded to CPU (`.device` == CPU), but its intended
     # compute device is the accelerator. The encode must run on the intended compute device, not the current
     # residency, or the whole encode silently runs on the CPU.
     compute_device = torch.device("meta")
-    text_encoder = _run_encode(monkeypatch, compute_device)
+    text_encoder, _ = _run_encode(monkeypatch, compute_device)
     assert text_encoder.forward_input_device == compute_device
 
 
 def test_anima_qwen3_encode_uses_cpu_for_cpu_only_model(monkeypatch):
     # A cpu_only encoder has compute_device == CPU; the encode must run on the CPU.
-    text_encoder = _run_encode(monkeypatch, torch.device("cpu"))
+    text_encoder, _ = _run_encode(monkeypatch, torch.device("cpu"))
     assert text_encoder.forward_input_device == torch.device("cpu")
+
+
+def test_parse_anima_prompt_attention() -> None:
+    prompt, ranges = parse_prompt_attention(r"a (red) fox, (long hair:2), \(literal\), ((soft))")
+
+    assert prompt == "a red fox, long hair, (literal), soft"
+    assert [(prompt[start:end], weight) for start, end, weight in ranges] == [
+        ("a ", 1.0),
+        ("red", 1.1),
+        (" fox, ", 1.0),
+        ("long hair", 2.0),
+        (", (literal), ", 1.0),
+        ("soft", 1.1**2),
+    ]
+
+
+def test_anima_weighted_t5_tokenization_preserves_clean_prompt_ids() -> None:
+    tokenizer = load_bundled_t5_tokenizer()
+    plain_prompt, plain_ranges = parse_prompt_attention("a portrait, long hair, blue eyes")
+    weighted_prompt, weighted_ranges = parse_prompt_attention("a portrait, (long hair:2), blue eyes")
+
+    plain_ids, plain_weights = tokenize_t5_with_weights(tokenizer, plain_prompt, plain_ranges, max_length=512)
+    weighted_ids, weighted_weights = tokenize_t5_with_weights(
+        tokenizer, weighted_prompt, weighted_ranges, max_length=512
+    )
+
+    assert plain_prompt == weighted_prompt
+    torch.testing.assert_close(plain_ids, weighted_ids)
+    torch.testing.assert_close(plain_weights, torch.ones_like(plain_weights))
+
+    offsets = tokenizer(weighted_prompt, return_offsets_mapping=True).offset_mapping
+    long_hair_start = weighted_prompt.index("long hair")
+    long_hair_end = long_hair_start + len("long hair")
+    expected_weights = torch.tensor(
+        [
+            2.0 if token_end > long_hair_start and token_start < long_hair_end else 1.0
+            for token_start, token_end in offsets
+        ]
+    )
+    torch.testing.assert_close(weighted_weights, expected_weights)
+
+
+def test_anima_explicit_weight_one_matches_unweighted_prompt() -> None:
+    tokenizer = load_bundled_t5_tokenizer()
+    plain_prompt, plain_ranges = parse_prompt_attention("long hair")
+    weighted_prompt, weighted_ranges = parse_prompt_attention("(long hair:1)")
+
+    plain_ids, plain_weights = tokenize_t5_with_weights(tokenizer, plain_prompt, plain_ranges, max_length=512)
+    weighted_ids, weighted_weights = tokenize_t5_with_weights(
+        tokenizer, weighted_prompt, weighted_ranges, max_length=512
+    )
+
+    assert plain_prompt == weighted_prompt
+    torch.testing.assert_close(plain_ids, weighted_ids)
+    torch.testing.assert_close(plain_weights, weighted_weights)
+
+
+def test_anima_qwen3_does_not_receive_weighting_markup(monkeypatch) -> None:
+    _, tokenizer = _run_encode(monkeypatch, torch.device("cpu"), "a portrait, (long hair:2), blue eyes")
+
+    assert tokenizer.prompt == "a portrait, long hair, blue eyes"