Enferlain-ComfyUI-A1111-cond/api/tokenize.py

"""
Token Counter API for A1111 Prompt Node

Provides a REST endpoint for live tokenization feedback.
Uses word-by-word tokenization with manual position tracking.
"""

from aiohttp import web
import server
import re

# Lazy-loaded tokenizer instance
_tokenizer = None


def get_tokenizer():
    """Get or create the SD tokenizer."""
    global _tokenizer
    if _tokenizer is None:
        from comfy.sd1_clip import SDTokenizer

        _tokenizer = SDTokenizer()
    return _tokenizer


def strip_a1111_syntax(text: str) -> str:
    r"""
    Strip A1111 emphasis/scheduling syntax from text, leaving only tokenizable content.

    This matches what ComfyUI's clip.tokenize() does internally when it parses
    emphasis syntax like (word:1.2) - only the word gets tokenized, not the
    parentheses or weight numbers.

    For alternation/scheduling, keeps the LONGEST option to show worst-case token count.

    Handles:
    - Emphasis: (word:1.2) -> word, (word) -> word
    - Negative emphasis: [word] -> word, [word:0.5] -> word
    - Scheduling: [from:to:when] -> max(from, to) (keeps longer one)
    - Alternation: [A|B|C] -> max(A, B, C) (keeps longest option)
    - Escaped chars: \( \) \[ \] -> ( ) [ ]
    """
    # First, handle escaped characters - replace with placeholders
    text = text.replace("\\(", "\x00LPAREN\x00")
    text = text.replace("\\)", "\x00RPAREN\x00")
    text = text.replace("\\[", "\x00LBRACK\x00")
    text = text.replace("\\]", "\x00RBRACK\x00")

    # Handle bracket expressions: [A|B|C] or [from:to:when]
    # Keep the LONGEST option for worst-case token counting
    def keep_longest_option(match):
        content = match.group(1)

        if "|" in content:
            # It's alternation [A|B|C] - find longest option
            options = content.split("|")
            # Remove any trailing :number from last option (scheduled alternation)
            # e.g., [A|B:0.5] -> options = ["A", "B:0.5"]
            last = options[-1]
            if ":" in last:
                # Check if it ends with :number
                colon_match = re.match(r"^(.+?)::?[\d.]+$", last)
                if colon_match:
                    options[-1] = colon_match.group(1)

            # Return the longest option
            return max(options, key=len)

        elif ":" in content:
            # It's scheduling [from:to:when] - keep longer of from/to
            parts = content.split(":")
            if len(parts) >= 2:
                # parts[0] = from, parts[1] = to, parts[2+] = when (numbers)
                from_part = parts[0]
                to_part = parts[1] if len(parts) > 1 else ""
                # Return the longer of from/to
                return from_part if len(from_part) >= len(to_part) else to_part
            return content

        else:
            # Simple bracket emphasis [word] - just return content
            return content

    # Process bracket expressions - non-greedy match for innermost brackets
    # Repeat to handle nested structures
    prev_text = None
    while prev_text != text:
        prev_text = text
        text = re.sub(r"\[([^\[\]]*)\]", keep_longest_option, text)

    # Remove weight specifications: :1.2) at end of emphasis
    # This handles (word:1.2)
    text = re.sub(r":[\d.]+(?=\))", "", text)

    # Remove any remaining weight specs that might be floating
    text = re.sub(r":[\d.]+(?=\s|$)", "", text)

    # Remove parentheses (they're just syntax markers now)
    text = text.replace("(", " ")
    text = text.replace(")", " ")

    # Restore escaped characters as literal chars
    text = text.replace("\x00LPAREN\x00", "(")
    text = text.replace("\x00RPAREN\x00", ")")
    text = text.replace("\x00LBRACK\x00", "[")
    text = text.replace("\x00RBRACK\x00", "]")

    # Clean up multiple spaces
    text = re.sub(r"\s+", " ", text)

    return text.strip()


@server.PromptServer.instance.routes.post("/a1111_prompt/tokenize")
async def tokenize_prompt(request):
    """
    API endpoint for live tokenization.

    Returns token count per 77-token sequence and character positions
    where boundaries fall using word-by-word tokenization.
    BREAK forces a new sequence (matching A1111/parser behavior).
    """
    data = await request.json()
    text = data.get("text", "")

    try:
        tokenizer = get_tokenizer()
        hf_tokenizer = tokenizer.tokenizer  # Access underlying HuggingFace tokenizer

        # Split by BREAK first (matching parser.py behavior)
        break_pattern = r"\s*\bBREAK\b\s*"
        break_matches = list(re.finditer(break_pattern, text))
        break_segments = re.split(break_pattern, text)

        sequences = []
        boundaries = []
        tokens_detail = []  # Per-token info: [{text, id, chunk_idx}, ...]

        # Track position in original text
        current_text_offset = 0
        current_chunk_idx = 0

        for seg_idx, segment in enumerate(break_segments):
            # Find where this segment starts in the original text
            if segment:
                segment_start = text.find(segment, current_text_offset)
            else:
                segment_start = current_text_offset

            segment_text = segment.strip()

            if not segment_text:
                sequences.append(0)
                current_chunk_idx += 1
            else:
                # FIRST: Strip A1111 syntax from entire segment
                # This handles bracket expressions with spaces like [A|B C|D]
                clean_segment = strip_a1111_syntax(segment_text)

                if not clean_segment:
                    sequences.append(0)
                    current_chunk_idx += 1
                    current_text_offset = segment_start + len(segment)
                    continue

                # Word-by-word tokenization with position tracking
                # Split into words from the CLEANED segment
                word_pattern = r"(\S+)"
                words_with_pos = []

                for match in re.finditer(word_pattern, clean_segment):
                    word = match.group(1)
                    # Position relative to segment start
                    word_start = match.start()
                    word_end = match.end()
                    words_with_pos.append((word, word_start, word_end))

                # Tokenize each word and track cumulative token count
                chunk_size = 75
                current_chunk_tokens = 0
                chunk_sequences = []

                for word, word_start_rel, word_end_rel in words_with_pos:
                    # Words are already from the cleaned segment, just tokenize directly
                    word_tokens = hf_tokenizer.encode(word, add_special_tokens=False)
                    word_token_count = len(word_tokens)

                    # Check if adding this word would exceed chunk size
                    if (
                        current_chunk_tokens + word_token_count > chunk_size
                        and current_chunk_tokens > 0
                    ):
                        # Save current chunk
                        chunk_sequences.append(current_chunk_tokens)

                        # Add boundary marker at the END of the previous word
                        # (which is the start of this word)
                        boundary_char = segment_start + word_start_rel
                        boundaries.append({"char_pos": boundary_char, "type": "chunk"})

                        current_chunk_tokens = 0
                        current_chunk_idx += 1

                    # Decode each token to get the text representation
                    for token_id in word_tokens:
                        token_text = hf_tokenizer.decode([token_id])
                        tokens_detail.append(
                            {
                                "text": token_text,
                                "id": token_id,
                                "chunk": current_chunk_idx,
                            }
                        )

                    current_chunk_tokens += word_token_count

                # Don't forget the last chunk
                if current_chunk_tokens > 0:
                    chunk_sequences.append(current_chunk_tokens)

                sequences.extend(chunk_sequences if chunk_sequences else [0])
                current_chunk_idx += 1

            # Move past this segment
            current_text_offset = segment_start + len(segment)

            # Add BREAK boundary if not the last segment
            if seg_idx < len(break_matches):
                boundaries.append(
                    {"char_pos": break_matches[seg_idx].start(), "type": "break"}
                )
                # Add BREAK marker to tokens
                tokens_detail.append(
                    {
                        "text": "BREAK",
                        "id": None,
                        "chunk": current_chunk_idx - 1,
                        "is_break": True,
                    }
                )

        if not sequences:
            sequences = [0]

        # Calculate stats
        word_count = len(text.split())
        char_count = len(text)

        return web.json_response(
            {
                "sequences": sequences,
                "boundaries": boundaries,
                "tokens": tokens_detail,
                "stats": {
                    "total_tokens": sum(sequences),
                    "chunks": len(sequences),
                    "words": word_count,
                    "characters": char_count,
                },
            }
        )
    except Exception as e:
        import traceback

        traceback.print_exc()
        return web.json_response(
            {"error": str(e), "sequences": None, "boundaries": None}, status=500
        )