428 lines
18 KiB
Python
428 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Mapping, Sequence
|
|
|
|
import torch
|
|
|
|
from .types import V9_VIRTUAL_TOKEN_COUNT
|
|
|
|
|
|
class TriggerBindingError(ValueError):
|
|
pass
|
|
|
|
|
|
class TriggerPlaceholderError(TriggerBindingError):
|
|
pass
|
|
|
|
|
|
class TriggerLiteralNotFoundError(TriggerBindingError):
|
|
pass
|
|
|
|
|
|
class TriggerConflictError(TriggerBindingError):
|
|
pass
|
|
|
|
|
|
class TriggerTokenizerError(TriggerBindingError):
|
|
pass
|
|
|
|
|
|
class TriggerTokenizerParityError(TriggerTokenizerError):
|
|
pass
|
|
|
|
|
|
class TriggerTruncationError(TriggerBindingError):
|
|
pass
|
|
|
|
|
|
class TriggerAtomicityError(TriggerBindingError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class TriggerBindingMetadata:
|
|
raw_text: str
|
|
resolved_text: str
|
|
rendered_text: str
|
|
literal: str
|
|
character_spans: tuple[tuple[int, int], ...]
|
|
token_spans: tuple[tuple[int, int], ...]
|
|
token_indices: tuple[int, ...]
|
|
virtual_token_indices: tuple[int, ...]
|
|
occurrence_indices: tuple[int, ...]
|
|
input_ids: tuple[int, ...]
|
|
attention_mask: tuple[int, ...]
|
|
trigger_mask: tuple[int, ...]
|
|
atomic_token_id: int | None
|
|
lookup_token_id: int | None
|
|
|
|
@property
|
|
def occurrence_count(self) -> int:
|
|
return len(self.character_spans)
|
|
|
|
@property
|
|
def slot_count(self) -> int:
|
|
return len(self.token_indices)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class TriggerBindingBatch:
|
|
items: tuple[TriggerBindingMetadata, ...]
|
|
input_ids: torch.Tensor
|
|
attention_mask: torch.Tensor
|
|
trigger_mask: torch.Tensor
|
|
virtual_token_indices: torch.Tensor
|
|
occurrence_indices: torch.Tensor
|
|
|
|
|
|
def resolve_trigger_literal(raw_text: str, literal: str, placeholder: str = "[trigger]", require_placeholder: bool = True, reject_literal_conflicts: bool = True) -> tuple[str, tuple[tuple[int, int], ...]]:
|
|
if not isinstance(raw_text, str) or not isinstance(literal, str) or not literal or not placeholder:
|
|
raise TriggerPlaceholderError("raw text, placeholder and literal must be valid non-empty strings")
|
|
if require_placeholder and placeholder not in raw_text:
|
|
raise TriggerPlaceholderError(f"caption does not contain required placeholder {placeholder!r}")
|
|
if reject_literal_conflicts and literal in raw_text:
|
|
raise TriggerConflictError("raw caption already contains the literal trigger")
|
|
parts = raw_text.split(placeholder)
|
|
output: list[str] = []
|
|
spans: list[tuple[int, int]] = []
|
|
cursor = 0
|
|
for index, part in enumerate(parts):
|
|
output.append(part)
|
|
cursor += len(part)
|
|
if index < len(parts) - 1:
|
|
spans.append((cursor, cursor + len(literal)))
|
|
output.append(literal)
|
|
cursor += len(literal)
|
|
return "".join(output), tuple(spans)
|
|
|
|
|
|
def find_literal_spans(text: str, literal: str) -> tuple[tuple[int, int], ...]:
|
|
spans: list[tuple[int, int]] = []
|
|
start = 0
|
|
while literal:
|
|
index = text.find(literal, start)
|
|
if index < 0:
|
|
break
|
|
spans.append((index, index + len(literal)))
|
|
start = index + len(literal)
|
|
return tuple(spans)
|
|
|
|
|
|
def render_chat_prompt(tokenizer: Any, text: str, add_generation_prompt: bool = True) -> str:
|
|
for content in ([{"type": "text", "text": text}], text):
|
|
try:
|
|
rendered = tokenizer.apply_chat_template(
|
|
[{"role": "user", "content": content}],
|
|
add_generation_prompt=add_generation_prompt,
|
|
tokenize=False,
|
|
)
|
|
if isinstance(rendered, str):
|
|
return rendered
|
|
except (AttributeError, TypeError, ValueError):
|
|
continue
|
|
raise TriggerTokenizerError("tokenizer must support the Ideogram4 Qwen chat template")
|
|
|
|
|
|
def _flat(value: Any, name: str) -> list[Any]:
|
|
if isinstance(value, torch.Tensor):
|
|
value = value.detach().cpu().tolist()
|
|
if isinstance(value, tuple):
|
|
value = list(value)
|
|
if value and isinstance(value[0], (list, tuple)):
|
|
if len(value) != 1:
|
|
raise TriggerTokenizerError(f"{name} unexpectedly returned a batch")
|
|
value = list(value[0])
|
|
if not isinstance(value, list):
|
|
raise TriggerTokenizerError(f"tokenizer output {name!r} must be a sequence")
|
|
return value
|
|
|
|
|
|
def _tokenize_ids(tokenizer: Any, text: str) -> tuple[int, ...]:
|
|
encoded = tokenizer(text, add_special_tokens=False, truncation=False)
|
|
return tuple(int(value) for value in _flat(encoded["input_ids"], "input_ids"))
|
|
|
|
|
|
def _tokenize_offsets(
|
|
tokenizer: Any,
|
|
text: str,
|
|
max_length: int | None = None,
|
|
) -> tuple[list[int], list[int], list[tuple[int, int]]]:
|
|
if getattr(tokenizer, "is_fast", None) is not True:
|
|
raise TriggerTokenizerError("trigger binding requires a fast tokenizer")
|
|
kwargs: dict[str, Any] = {
|
|
"add_special_tokens": False,
|
|
"return_offsets_mapping": True,
|
|
"truncation": max_length is not None,
|
|
}
|
|
if max_length is not None:
|
|
if max_length <= 0:
|
|
raise TriggerTruncationError("max_length cannot fit the configured virtual trigger tokens")
|
|
kwargs["max_length"] = int(max_length)
|
|
encoded = tokenizer(text, **kwargs)
|
|
ids = [int(value) for value in _flat(encoded["input_ids"], "input_ids")]
|
|
masks = [int(value) for value in _flat(encoded.get("attention_mask", [1] * len(ids)), "attention_mask")]
|
|
raw_offsets = encoded["offset_mapping"]
|
|
if isinstance(raw_offsets, torch.Tensor):
|
|
raw_offsets = raw_offsets.detach().cpu().tolist()
|
|
if isinstance(raw_offsets, tuple):
|
|
raw_offsets = list(raw_offsets)
|
|
if (
|
|
isinstance(raw_offsets, list)
|
|
and len(raw_offsets) == 1
|
|
and isinstance(raw_offsets[0], list)
|
|
and (not raw_offsets[0] or isinstance(raw_offsets[0][0], (list, tuple)))
|
|
):
|
|
raw_offsets = raw_offsets[0]
|
|
if not isinstance(raw_offsets, list):
|
|
raise TriggerTokenizerError("offset_mapping must be a sequence")
|
|
offsets = [tuple(map(int, pair)) for pair in raw_offsets]
|
|
if not (len(ids) == len(masks) == len(offsets)):
|
|
raise TriggerTokenizerError("tokenizer output lengths differ")
|
|
return ids, masks, offsets
|
|
|
|
|
|
def _tokenizer_children(value: Any) -> list[Any]:
|
|
children: list[Any] = []
|
|
for attribute in (
|
|
"tokenizer",
|
|
"clip",
|
|
"clip_l",
|
|
"qwen3_vl",
|
|
"qwen3vl_8b",
|
|
"qwen3_8b",
|
|
):
|
|
children.append(getattr(value, attribute, None))
|
|
dynamic_names = (
|
|
getattr(value, "clip", None),
|
|
getattr(value, "clip_name", None),
|
|
)
|
|
for name in dynamic_names:
|
|
if isinstance(name, str) and name:
|
|
children.append(getattr(value, name, None))
|
|
return children
|
|
|
|
|
|
def resolve_huggingface_tokenizer(tokenizer_or_clip: Any) -> Any:
|
|
pending, seen = [tokenizer_or_clip], set()
|
|
while pending:
|
|
value = pending.pop(0)
|
|
if value is None or id(value) in seen:
|
|
continue
|
|
seen.add(id(value))
|
|
if hasattr(value, "get_vocab") and callable(value) and hasattr(value, "apply_chat_template"):
|
|
return value
|
|
pending.extend(_tokenizer_children(value))
|
|
raise TriggerTokenizerError(
|
|
"could not locate the Hugging Face tokenizer inside the connected ComfyUI tokenizer wrapper"
|
|
)
|
|
|
|
|
|
def _tokenizer_resource_path(tokenizer: Any) -> Path:
|
|
candidates = [getattr(tokenizer, "name_or_path", None)]
|
|
init_kwargs = getattr(tokenizer, "init_kwargs", {})
|
|
if isinstance(init_kwargs, Mapping):
|
|
candidates.append(init_kwargs.get("name_or_path"))
|
|
for candidate in candidates:
|
|
if candidate and Path(candidate).expanduser().is_dir():
|
|
return Path(candidate).expanduser().resolve()
|
|
raise TriggerTokenizerError("connected tokenizer does not expose a local resource directory")
|
|
|
|
|
|
def _special_snapshot(tokenizer: Any) -> tuple[Any, ...]:
|
|
return (
|
|
tuple(sorted((str(k), str(v)) for k, v in dict(getattr(tokenizer, "special_tokens_map", {}) or {}).items())),
|
|
tuple(int(v) for v in getattr(tokenizer, "all_special_ids", ()) or ()),
|
|
)
|
|
|
|
|
|
def validate_fast_slow_tokenizer_parity(slow_tokenizer: Any, fast_tokenizer: Any, rendered_text: str, literal: str | None = None) -> None:
|
|
if slow_tokenizer.get_vocab() != fast_tokenizer.get_vocab():
|
|
raise TriggerTokenizerParityError("fast and slow tokenizer vocabularies differ")
|
|
if _special_snapshot(slow_tokenizer) != _special_snapshot(fast_tokenizer):
|
|
raise TriggerTokenizerParityError("fast and slow tokenizer special-token configuration differs")
|
|
parity_text = rendered_text.replace(literal, "") if literal else rendered_text
|
|
if _tokenize_ids(slow_tokenizer, parity_text) != _tokenize_ids(fast_tokenizer, parity_text):
|
|
raise TriggerTokenizerParityError("fast and slow tokenizer IDs differ")
|
|
|
|
|
|
def create_private_ideogram4_fast_tokenizer(connected_tokenizer_or_clip: Any, parity_text: str = "Ideogram4 tokenizer parity check.", literal: str | None = None, validate_parity: bool = True) -> Any:
|
|
slow = resolve_huggingface_tokenizer(connected_tokenizer_or_clip)
|
|
try:
|
|
from transformers import Qwen2TokenizerFast
|
|
except ImportError as exc:
|
|
raise TriggerTokenizerError("transformers.Qwen2TokenizerFast is required") from exc
|
|
try:
|
|
fast = Qwen2TokenizerFast.from_pretrained(str(_tokenizer_resource_path(slow)), local_files_only=True)
|
|
except (OSError, ValueError) as exc:
|
|
raise TriggerTokenizerError("failed to create private local Qwen2TokenizerFast") from exc
|
|
if validate_parity:
|
|
validate_fast_slow_tokenizer_parity(slow, fast, parity_text, literal)
|
|
return fast
|
|
|
|
|
|
def register_atomic_literal(
|
|
tokenizer: Any,
|
|
literal: str,
|
|
vocab_limit: int,
|
|
) -> tuple[int, int]:
|
|
lookup_ids = getattr(tokenizer, "_gen2_v9_lookup_token_ids", None)
|
|
if lookup_ids is None:
|
|
lookup_ids = {}
|
|
try:
|
|
setattr(tokenizer, "_gen2_v9_lookup_token_ids", lookup_ids)
|
|
except (AttributeError, TypeError) as exc:
|
|
raise TriggerTokenizerError("private tokenizer cannot retain V9 literal lookup metadata") from exc
|
|
lookup_token_id = lookup_ids.get(literal)
|
|
if lookup_token_id is None:
|
|
source_ids = _tokenize_ids(tokenizer, literal)
|
|
if not source_ids or any(token_id < 0 or token_id >= vocab_limit for token_id in source_ids):
|
|
raise TriggerAtomicityError(
|
|
"literal pre-registration token IDs must remain inside the frozen Qwen embedding table"
|
|
)
|
|
lookup_token_id = int(source_ids[0])
|
|
lookup_ids[literal] = lookup_token_id
|
|
added = tokenizer.add_tokens([literal], special_tokens=True)
|
|
if added not in (0, 1):
|
|
raise TriggerTokenizerError(f"unexpected number of registered tokens: {added}")
|
|
ids = _tokenize_ids(tokenizer, literal)
|
|
if len(ids) != 1:
|
|
raise TriggerAtomicityError(f"literal must map to exactly one token ID, got {list(ids)}")
|
|
atomic_token_id = ids[0]
|
|
if getattr(tokenizer, "unk_token_id", None) == atomic_token_id:
|
|
raise TriggerAtomicityError("literal maps to unknown token ID")
|
|
return atomic_token_id, lookup_token_id
|
|
|
|
|
|
def _overlapping_token(offsets: Sequence[tuple[int, int]], span: tuple[int, int]) -> int:
|
|
start, end = span
|
|
matches = [index for index, (token_start, token_end) in enumerate(offsets) if token_end > token_start and token_start < end and token_end > start]
|
|
if len(matches) != 1:
|
|
raise TriggerAtomicityError("registered literal must occupy one contextual token")
|
|
index = matches[0]
|
|
if offsets[index][0] > start or offsets[index][1] < end:
|
|
raise TriggerTruncationError("trigger occurrence was partially tokenized")
|
|
return index
|
|
|
|
|
|
def bind_trigger_prompt(
|
|
tokenizer: Any,
|
|
raw_text: str,
|
|
literal: str,
|
|
vocab_limit: int | None = None,
|
|
max_length: int | None = None,
|
|
stock_literal: bool = False,
|
|
) -> TriggerBindingMetadata:
|
|
if not isinstance(raw_text, str) or not isinstance(literal, str) or not literal:
|
|
raise TriggerLiteralNotFoundError("text and literal must be valid, and literal must not be empty")
|
|
source_spans = find_literal_spans(raw_text, literal)
|
|
if not source_spans:
|
|
raise TriggerLiteralNotFoundError(f"text does not contain the configured literal trigger {literal!r}")
|
|
rendered = render_chat_prompt(tokenizer, raw_text)
|
|
rendered_spans = find_literal_spans(rendered, literal)
|
|
if len(rendered_spans) != len(source_spans):
|
|
raise TriggerConflictError(
|
|
"chat template changed or duplicated literal-trigger occurrences; offset mapping is ambiguous"
|
|
)
|
|
if stock_literal:
|
|
ids, attention, offsets = _tokenize_offsets(tokenizer, rendered, max_length)
|
|
for span in rendered_spans:
|
|
matches = [
|
|
index
|
|
for index, (token_start, token_end) in enumerate(offsets)
|
|
if token_end > token_start and token_start < span[1] and token_end > span[0]
|
|
]
|
|
if not matches or offsets[matches[0]][0] > span[0] or offsets[matches[-1]][1] < span[1]:
|
|
raise TriggerTruncationError("stock literal occurrence was truncated or not tokenized")
|
|
return TriggerBindingMetadata(raw_text, raw_text, rendered, literal, rendered_spans, (), (), (), (), tuple(ids), tuple(attention), tuple(0 for _ in ids), None, None)
|
|
if vocab_limit is None:
|
|
vocab_limit = len(tokenizer.get_vocab())
|
|
atomic_id, lookup_id = register_atomic_literal(tokenizer, literal, vocab_limit)
|
|
expansion_slots = len(rendered_spans) * (V9_VIRTUAL_TOKEN_COUNT - 1)
|
|
tokenizer_max_length = None if max_length is None else max_length - expansion_slots
|
|
if tokenizer_max_length is not None and tokenizer_max_length < len(rendered_spans):
|
|
raise TriggerTruncationError("max_length cannot fit the configured virtual trigger tokens")
|
|
ids, attention, offsets = _tokenize_offsets(tokenizer, rendered, tokenizer_max_length)
|
|
occurrence_positions = [_overlapping_token(offsets, span) for span in rendered_spans]
|
|
if len(set(occurrence_positions)) != len(occurrence_positions):
|
|
raise TriggerAtomicityError("multiple literal occurrences resolved to the same contextual token")
|
|
expanded_ids: list[int] = []
|
|
expanded_attention: list[int] = []
|
|
trigger_mask: list[int] = []
|
|
virtual_indices: list[int] = []
|
|
occurrence_indices: list[int] = []
|
|
token_spans: list[tuple[int, int]] = []
|
|
occurrence_by_position = {position: index for index, position in enumerate(occurrence_positions)}
|
|
for position, (token_id, mask) in enumerate(zip(ids, attention)):
|
|
occurrence = occurrence_by_position.get(position)
|
|
if occurrence is None:
|
|
expanded_ids.append(token_id)
|
|
expanded_attention.append(mask)
|
|
trigger_mask.append(0)
|
|
virtual_indices.append(-1)
|
|
occurrence_indices.append(-1)
|
|
continue
|
|
if token_id != atomic_id:
|
|
raise TriggerAtomicityError("literal contextual token ID differs from registered atomic ID")
|
|
start = len(expanded_ids)
|
|
for virtual_index in range(V9_VIRTUAL_TOKEN_COUNT):
|
|
expanded_ids.append(atomic_id)
|
|
expanded_attention.append(mask)
|
|
trigger_mask.append(1)
|
|
virtual_indices.append(virtual_index)
|
|
occurrence_indices.append(occurrence)
|
|
token_spans.append((start, start + V9_VIRTUAL_TOKEN_COUNT))
|
|
if max_length is not None and len(expanded_ids) > max_length:
|
|
raise TriggerTruncationError(
|
|
f"four-slot expansion requires {len(expanded_ids)} tokens, exceeding max_length={max_length}"
|
|
)
|
|
token_indices = tuple(index for index, value in enumerate(trigger_mask) if value)
|
|
return TriggerBindingMetadata(
|
|
raw_text, raw_text, rendered, literal, rendered_spans, tuple(token_spans), token_indices,
|
|
tuple(virtual_indices[index] for index in token_indices),
|
|
tuple(occurrence_indices[index] for index in token_indices), tuple(expanded_ids),
|
|
tuple(expanded_attention), tuple(trigger_mask), atomic_id, lookup_id,
|
|
)
|
|
|
|
|
|
def bind_trigger_batch(tokenizer: Any, raw_texts: Sequence[str], literal: str, pad_token_id: int | None = None, **kwargs: Any) -> TriggerBindingBatch:
|
|
if not raw_texts:
|
|
raise TriggerBindingError("trigger binding batch must not be empty")
|
|
items = tuple(bind_trigger_prompt(tokenizer, text, literal, **kwargs) for text in raw_texts)
|
|
length = max(len(item.input_ids) for item in items)
|
|
pad = int(pad_token_id if pad_token_id is not None else getattr(tokenizer, "pad_token_id", 0) or 0)
|
|
def rows(name: str, fill: int) -> list[list[int]]:
|
|
return [list(getattr(item, name)) + [fill] * (length - len(getattr(item, name))) for item in items]
|
|
|
|
virtual_rows: list[list[int]] = []
|
|
occurrence_rows: list[list[int]] = []
|
|
for item in items:
|
|
virtual = [-1] * len(item.input_ids)
|
|
occurrences = [-1] * len(item.input_ids)
|
|
for token_index, virtual_index, occurrence_index in zip(
|
|
item.token_indices, item.virtual_token_indices, item.occurrence_indices
|
|
):
|
|
virtual[token_index] = virtual_index
|
|
occurrences[token_index] = occurrence_index
|
|
virtual_rows.append(virtual + [-1] * (length - len(virtual)))
|
|
occurrence_rows.append(occurrences + [-1] * (length - len(occurrences)))
|
|
return TriggerBindingBatch(
|
|
items,
|
|
torch.tensor(rows("input_ids", pad), dtype=torch.long),
|
|
torch.tensor(rows("attention_mask", 0), dtype=torch.long),
|
|
torch.tensor(rows("trigger_mask", 0), dtype=torch.bool),
|
|
torch.tensor(virtual_rows, dtype=torch.long),
|
|
torch.tensor(occurrence_rows, dtype=torch.long),
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"TriggerAtomicityError", "TriggerBindingBatch", "TriggerBindingError", "TriggerBindingMetadata",
|
|
"TriggerConflictError", "TriggerLiteralNotFoundError", "TriggerPlaceholderError", "TriggerTokenizerError", "TriggerTokenizerParityError",
|
|
"TriggerTruncationError", "bind_trigger_batch", "bind_trigger_prompt",
|
|
"create_private_ideogram4_fast_tokenizer", "find_literal_spans", "register_atomic_literal",
|
|
"resolve_huggingface_tokenizer",
|
|
"render_chat_prompt", "resolve_trigger_literal", "validate_fast_slow_tokenizer_parity",
|
|
]
|