ComfyUI-Flux2Klein-Enhancer/flux2_sectioned_encoder.py

import re
from typing import Dict, List, Optional, Tuple

import torch


# Klein's chat template (matches ComfyUI's KleinTokenizer.llama_template).
_KLEIN_CHAT_TEMPLATE = (
    "<|im_start|>user\n{}<|im_end|>\n"
    "<|im_start|>assistant\n<think>\n\n</think>\n\n"
)


def _get_hf_tokenizer(clip):
    """Return the underlying HF tokenizer (Qwen3-8B or Qwen3-4B variant)."""
    tok = getattr(clip, "tokenizer", None)
    if tok is None:
        return None
    for attr in ("qwen3_8b", "qwen3_4b"):
        sub = getattr(tok, attr, None)
        if sub is not None and hasattr(sub, "tokenizer"):
            return sub.tokenizer
    return None


def _count_tokens(hf_tok, text: str) -> int:
    if not text:
        return 0
    out = hf_tok(text, add_special_tokens=False, return_tensors=None)
    return len(out["input_ids"])


def _compute_wrapper_lengths(hf_tok) -> Tuple[int, int]:
    """Return (prefix_len, suffix_len) of Klein's chat-template wrapper."""
    prefix, suffix = _KLEIN_CHAT_TEMPLATE.split("{}")
    return _count_tokens(hf_tok, prefix), _count_tokens(hf_tok, suffix)


def _compute_section_ranges(
    hf_tok,
    sections: Dict[str, str],
    separator: str,
) -> Optional[Dict[str, Tuple[int, int]]]:
    """Return {'front': (s, e), 'mid': (s, e), 'end': (s, e)} in encoded-token positions.

    None if the tokenizer is missing — caller falls back to encoding without metadata.
    """
    if hf_tok is None:
        return None

    prefix_len, _ = _compute_wrapper_lengths(hf_tok)
    sep_len = _count_tokens(hf_tok, separator) if separator else 0

    # Tokenize each section's bare text.
    front_n = _count_tokens(hf_tok, sections.get("front", ""))
    mid_n = _count_tokens(hf_tok, sections.get("mid", ""))
    end_n = _count_tokens(hf_tok, sections.get("end", ""))

    # Layout: prefix | front | (sep if both front+mid present) | mid | (sep if mid+end or front+end with no mid) | end
    pos = prefix_len
    ranges: Dict[str, Tuple[int, int]] = {}

    ranges["front"] = (pos, pos + front_n)
    pos += front_n
    if front_n > 0 and mid_n > 0:
        pos += sep_len

    ranges["mid"] = (pos, pos + mid_n)
    pos += mid_n
    if (mid_n > 0 and end_n > 0) or (front_n > 0 and mid_n == 0 and end_n > 0):
        pos += sep_len

    ranges["end"] = (pos, pos + end_n)
    return ranges


def _parse_marker_sections(text: str) -> Optional[Dict[str, str]]:
    """Parse [FRONT]/[MID]/[END] markers from a single combined prompt."""
    if not text:
        return None
    pattern = r"\[(FRONT|MID|END)\](.*?)(?=\[(?:FRONT|MID|END)\]|$)"
    matches = re.findall(pattern, text, re.DOTALL | re.IGNORECASE)
    if not matches:
        return None
    sections = {"front": "", "mid": "", "end": ""}
    for name, content in matches:
        sections[name.lower()] = content.strip()
    return sections


_SEPARATORS = {"comma": ", ", "period": ". ", "space": " ", "newline": "\n"}


class Flux2KleinSectionedEncoder:
    """Tokenize a 3-section prompt and emit per-section token ranges as
    conditioning metadata that the Detail Controller can read."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "clip": ("CLIP",),
            },
            "optional": {
                "front_text": ("STRING", {
                    "multiline": True, "default": "",
                    "tooltip": "FRONT section text.",
                }),
                "mid_text": ("STRING", {
                    "multiline": True, "default": "",
                    "tooltip": "MID section text.",
                }),
                "end_text": ("STRING", {
                    "multiline": True, "default": "",
                    "tooltip": "END section text.",
                }),
                "combined_prompt": ("STRING", {
                    "multiline": True, "default": "",
                    "tooltip": "Optional. Single prompt with [FRONT]/[MID]/[END] markers — overrides the three text boxes when non-empty and contains markers.",
                }),
                "separator": (list(_SEPARATORS.keys()), {
                    "default": "comma",
                    "tooltip": "How to join sections in the final prompt sent to Klein.",
                }),
                "show_preview": ("BOOLEAN", {"default": True}),
                "debug": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ("CONDITIONING", "STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = ("conditioning", "front_section", "mid_section", "end_section", "full_prompt")
    FUNCTION = "encode_sectioned"
    CATEGORY = "conditioning/flux2klein"
    OUTPUT_NODE = True

    def encode_sectioned(self, clip, front_text="", mid_text="", end_text="",
                         combined_prompt="", separator="comma",
                         show_preview=True, debug=False):

        # Resolve sections: marker mode if combined_prompt has [FRONT] etc., else 3-box.
        sections: Dict[str, str]
        marker_sections = _parse_marker_sections(combined_prompt)
        if marker_sections:
            sections = marker_sections
        else:
            sections = {
                "front": front_text or "",
                "mid": mid_text or "",
                "end": end_text or "",
            }

        sep_str = _SEPARATORS.get(separator, ", ")

        # Build the final prompt — drop empty sections from the join so we
        # don't end up with trailing/leading separators.
        parts = [sections[k] for k in ("front", "mid", "end") if sections[k]]
        full_prompt = sep_str.join(parts)

        # Standard CLIP tokenize + encode path.
        tokens = clip.tokenize(full_prompt)
        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)

        # Compute per-section token boundaries via the underlying HF tokenizer.
        hf_tok = _get_hf_tokenizer(clip)
        ranges = _compute_section_ranges(hf_tok, sections, sep_str)

        meta: Dict = {"pooled_output": pooled}
        if ranges is not None:
            meta["klein_sections"] = ranges
            if debug:
                print(f"[SectionedEncoder] computed klein_sections: {ranges}")
        else:
            print("[SectionedEncoder] WARNING: HF tokenizer not accessible on CLIP — "
                  "no klein_sections metadata emitted. Detail Controller will fall back "
                  "to fixed 25/50/25 slicing.")

        conditioning = [[cond, meta]]

        if show_preview or debug:
            self._print_preview(sections, full_prompt, ranges, separator, hf_tok)

        return (
            conditioning,
            sections["front"],
            sections["mid"],
            sections["end"],
            full_prompt,
        )

    def _print_preview(self, sections, full_prompt, ranges, separator, hf_tok):
        lines = ["", "=" * 70, "FLUX.2 Klein Sectioned Encoding (v2)", "=" * 70]
        lines.append(f"Separator: {separator!r}")

        if hf_tok is not None:
            front_n = _count_tokens(hf_tok, sections["front"])
            mid_n = _count_tokens(hf_tok, sections["mid"])
            end_n = _count_tokens(hf_tok, sections["end"])
            prefix_n, suffix_n = _compute_wrapper_lengths(hf_tok)
            lines.append(f"Section token counts (HF-tokenizer-exact, no padding):")
            lines.append(f"  FRONT: {front_n} tokens   '{sections['front']}'")
            lines.append(f"  MID:   {mid_n} tokens   '{sections['mid']}'")
            lines.append(f"  END:   {end_n} tokens   '{sections['end']}'")
            lines.append(f"Klein wrapper overhead: prefix={prefix_n} suffix={suffix_n} tokens")
        else:
            lines.append("(HF tokenizer not accessible — token counts unavailable)")

        if ranges:
            lines.append("Encoded-sequence section ranges (used by Detail Controller):")
            for name in ("front", "mid", "end"):
                s, e = ranges[name]
                lines.append(f"  {name.upper():5s}  tokens [{s}:{e})  span={e - s}")
            lines.append("Wire this conditioning into Detail Controller — it will scale these exact ranges.")
        else:
            lines.append("WARNING: no section ranges computed (HF tokenizer unavailable).")

        lines.append("-" * 70)
        lines.append(f"Final prompt: {full_prompt!r}")
        lines.append("=" * 70)
        print("\n".join(lines))


NODE_CLASS_MAPPINGS = {
    "Flux2KleinSectionedEncoder": Flux2KleinSectionedEncoder,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Flux2KleinSectionedEncoder": "FLUX.2 Klein Sectioned Encoder",
}