diff --git a/invokeai/app/invocations/anima_text_encoder.py b/invokeai/app/invocations/anima_text_encoder.py index ae79327f3ce..e7de917b0cd 100644 --- a/invokeai/app/invocations/anima_text_encoder.py +++ b/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,7 +107,7 @@ 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: @@ -114,9 +115,9 @@ def _encode_prompt( - 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.""" diff --git a/invokeai/backend/anima/prompt_weighting.py b/invokeai/backend/anima/prompt_weighting.py new file mode 100644 index 00000000000..164095acaa8 --- /dev/null +++ b/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) diff --git a/tests/app/invocations/test_anima_text_encoder.py b/tests/app/invocations/test_anima_text_encoder.py index 4d92c69ce04..7753b3c7116 100644 --- a/tests/app/invocations/test_anima_text_encoder.py +++ b/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,13 +95,13 @@ 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): @@ -97,11 +109,72 @@ def test_anima_qwen3_encode_uses_compute_device(monkeypatch): # 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"