anima_mlx/runtime/tokenizer.py

"""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()