File size: 6,076 Bytes
3bdc93d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | """Tokenizer runtime for Anima prompt conditioning."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Sequence
@dataclass(frozen=True)
class TokenizerPaths:
qwen3_06b: Path
t5xxl: Path
@classmethod
def from_comfy_root(cls, comfy_root: str | Path) -> "TokenizerPaths":
root = Path(comfy_root)
bundled_tokenizer_dir = root / "tokenizers"
if bundled_tokenizer_dir.exists():
return cls(
qwen3_06b=bundled_tokenizer_dir / "qwen25_tokenizer",
t5xxl=bundled_tokenizer_dir / "t5_tokenizer",
)
text_encoder_dir = root / "comfy/text_encoders"
return cls(
qwen3_06b=text_encoder_dir / "qwen25_tokenizer",
t5xxl=text_encoder_dir / "t5_tokenizer",
)
def as_dict(self, *, relative_to: str | Path | None = None) -> dict[str, str]:
return {
"qwen3_06b": _stringify_path(self.qwen3_06b, relative_to),
"t5xxl": _stringify_path(self.t5xxl, relative_to),
}
@dataclass(frozen=True)
class PromptTokens:
qwen3_06b_ids: tuple[int, ...]
qwen3_06b_weights: tuple[float, ...]
t5xxl_ids: tuple[int, ...]
t5xxl_weights: tuple[float, ...]
@classmethod
def from_pairs(
cls,
*,
qwen3_06b: Sequence[tuple[int, float]],
t5xxl: Sequence[tuple[int, float]],
) -> "PromptTokens":
return cls(
qwen3_06b_ids=tuple(int(token_id) for token_id, _ in qwen3_06b),
qwen3_06b_weights=tuple(float(weight) for _, weight in qwen3_06b),
t5xxl_ids=tuple(int(token_id) for token_id, _ in t5xxl),
t5xxl_weights=tuple(float(weight) for _, weight in t5xxl),
)
def as_dict(self) -> dict[str, list[int] | list[float]]:
return {
"qwen3_06b_ids": list(self.qwen3_06b_ids),
"qwen3_06b_weights": list(self.qwen3_06b_weights),
"t5xxl_ids": list(self.t5xxl_ids),
"t5xxl_weights": list(self.t5xxl_weights),
}
class AnimaTokenizer:
"""Lightweight reproduction of ComfyUI's AnimaTokenizer contract."""
def __init__(self, qwen3_06b: Any, t5xxl: Any, *, paths: TokenizerPaths | None = None) -> None:
self.qwen3_06b = qwen3_06b
self.t5xxl = t5xxl
self.paths = paths
@classmethod
def from_comfy_root(cls, comfy_root: str | Path) -> "AnimaTokenizer":
from transformers import Qwen2Tokenizer, T5TokenizerFast
paths = TokenizerPaths.from_comfy_root(comfy_root)
qwen3_06b = Qwen2Tokenizer.from_pretrained(paths.qwen3_06b)
t5xxl = T5TokenizerFast.from_pretrained(paths.t5xxl)
return cls(qwen3_06b, t5xxl, paths=paths)
def tokenize(self, text: str) -> PromptTokens:
qwen_pairs = simple_comfy_tokenize(
self.qwen3_06b,
text,
has_start_token=False,
has_end_token=False,
pad_to_max_length=False,
max_length=99_999_999,
min_length=1,
pad_token=151_643,
)
t5_pairs = simple_comfy_tokenize(
self.t5xxl,
text,
has_start_token=False,
has_end_token=False,
pad_to_max_length=False,
max_length=99_999_999,
min_length=1,
pad_token=None,
)
return PromptTokens.from_pairs(qwen3_06b=qwen_pairs, t5xxl=t5_pairs)
def tokenize_pair(self, positive: str, negative: str) -> dict[str, dict[str, list[int] | list[float]]]:
return {
"positive": self.tokenize(positive).as_dict(),
"negative": self.tokenize(negative).as_dict(),
}
def golden_payload(
self,
positive: str,
negative: str,
*,
relative_to: str | Path | None = None,
) -> dict[str, Any]:
payload: dict[str, Any] = {
"schema_version": 1,
"source": "lightweight reproduction of ComfyUI AnimaTokenizer settings over ComfyUI tokenizer files",
}
if self.paths is not None:
payload["tokenizer_paths"] = self.paths.as_dict(relative_to=relative_to)
payload.update(self.tokenize_pair(positive, negative))
return payload
def decode_qwen(self, token_ids: Sequence[int], **kwargs: Any) -> str:
return self.qwen3_06b.decode(list(token_ids), **kwargs)
def simple_comfy_tokenize(
tokenizer: Any,
text: str,
*,
has_start_token: bool,
has_end_token: bool,
pad_to_max_length: bool,
max_length: int,
min_length: int | None,
pad_token: int | None,
) -> list[tuple[int, float]]:
empty = tokenizer("")["input_ids"]
tokens_start = 0
start_token = None
end_token = None
if has_start_token:
if empty:
tokens_start = 1
start_token = empty[0]
if has_end_token and len(empty) > 1:
end_token = empty[1]
elif has_end_token and empty:
end_token = empty[0]
if pad_token is None:
pad_token = end_token if end_token is not None else 0
end = -1 if has_end_token else 999_999_999_999
ids = tokenizer(text)["input_ids"][tokens_start:end]
batch: list[tuple[int, float]] = []
if start_token is not None:
batch.append((int(start_token), 1.0))
batch.extend((int(token_id), 1.0) for token_id in ids)
if end_token is not None:
batch.append((int(end_token), 1.0))
if pad_to_max_length and len(batch) < max_length:
batch.extend((int(pad_token), 1.0) for _ in range(max_length - len(batch)))
if min_length is not None and len(batch) < min_length:
batch.extend((int(pad_token), 1.0) for _ in range(min_length - len(batch)))
return batch
def _stringify_path(path: Path, relative_to: str | Path | None) -> str:
if relative_to is not None:
try:
return path.relative_to(Path(relative_to)).as_posix()
except ValueError:
pass
return path.as_posix()
|