olivv-cs-ComfyUI-FunPack/modules/conditioning/prompt_markup/nodes.py

"""The two markup nodes. See the module docstring for the shape."""

import torch
import torch.nn.functional as F
from comfy_api.latest import io

from ..._core import control as control_mod, dit_hooks, log, patching, registry as registry_mod, schema as schema_mod
from . import markup as mk

MARKUP = io.Custom("FUNUI_MARKUP")
KEY = "funui_markup.prompt"
PHRASES = "funui_markup_phrases"   # the transformer_options key the phrase probe reads
SIMUL = "funui_simultaneous"        # ... and the one Step merging reads: one whole conditioning per option, not yet refined by the model


def _ask(capability, *args):
    """The first model module with an answer, or None."""
    for _spec, provider in registry_mod.current().providers(capability):
        got = provider(*args)
        if got is not None:
            return got
    return None


def _say(message):
    log.once(f"prompt_markup:{message}", log.ALERT, "FunUI Prompt Markup", message)


def _alternatives(settings):
    """-> (how the main prompt reads an [A|B], whether A and B are also read separately), from Alternating prompts."""
    spec = registry_mod.current().specs.get("prompt_alternatives")
    if spec is None or control_mod.all_off(settings):
        return "while", False
    clean, _problems = schema_mod.check_values(spec, (settings or {}).get("prompt_alternatives"))
    mode = clean.get("main")
    return (mode if mode in mk.MAIN_MODES else "while"), bool(clean.get("separate"))


class FunUIPromptMarkup(io.ComfyNode):
    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="FunUIPromptMarkup",
            display_name="FunUI Prompt Markup",
            category="FunUI/Conditioning",
            description="Take (word:1.5) and [a|b] out of a prompt, so the encoder sees clean text "
                        "and the markup can be applied properly.",
            inputs=[io.String.Input("text", multiline=True, default=""),
                    io.Custom("FUNUI_SETTINGS").Input("settings", optional=True)],
            outputs=[io.String.Output(display_name="clean"),
                     MARKUP.Output(display_name="markup"),
                     io.String.Output(display_name="status")],
        )

    @classmethod
    def execute(cls, text: str, settings=None) -> io.NodeOutput:
        main, separate = _alternatives(settings)
        parsed = mk.parse(text or "", main=main)
        parsed["separate"] = separate
        n = {"weighted": len(parsed["weighted"]), "alternative": len(parsed["variants"])}
        status = (", ".join(f"{v} {k}" for k, v in n.items() if v) or "no markup")
        return io.NodeOutput(parsed["clean"], parsed, status)


def _variant_rows(rows, clip, tokenizer, texts):
    """The prompt rows of each option's own prompt, brought to the length of `rows` (the prompt as encoded) --
    [1, n, D] each, or None when one of them cannot be placed."""
    start, end = rows
    out = []
    for text in texts:
        encoded = clip.encode_from_tokens_scheduled(clip.tokenize(text))
        cond, meta = encoded[0][0], encoded[0][1]
        n = len(tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)["offset_mapping"])
        placed = _ask("prompt_rows", meta, int(cond.shape[1]), n)
        if placed is None:
            return None
        got = cond[:, placed[0]:placed[1], :]
        if got.shape[1] != end - start:
            # Linear along the token axis: a smoothing, not a word alignment.
            got = F.interpolate(got.float().transpose(1, 2), size=end - start, mode="linear",
                                align_corners=True).transpose(1, 2).to(got.dtype)
        out.append(got)
    return out


class FunUIApplyPromptMarkup(io.ComfyNode):
    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="FunUIApplyPromptMarkup",
            display_name="FunUI Apply Prompt Markup",
            category="FunUI/Conditioning",
            description="Apply the markup FunUI Prompt Markup took out: weights on the model, and one "
                        "prompt per [a|b] option for Step merging to use.",
            inputs=[io.Model.Input("model"), io.Clip.Input("clip"),
                    io.Conditioning.Input("positive"), io.Latent.Input("latent"),
                    MARKUP.Input("markup")],
            outputs=[io.Model.Output(display_name="model"),
                     io.Conditioning.Output(display_name="positive"),
                     io.String.Output(display_name="status")],
        )

    @classmethod
    def execute(cls, model, clip, positive, latent, markup) -> io.NodeOutput:
        patched = patching.clone(model)
        patching.strip(patched, "funui_markup.")
        dit_hooks.publish(patched, SIMUL, None)
        weighted = (markup or {}).get("weighted")
        variants = (markup or {}).get("variants") if (markup or {}).get("separate") else None
        if not (weighted or variants):
            return io.NodeOutput(patched, positive, "no markup")

        def refuse(why):
            _say(f"markup NOT applied: {why}. The prompt still went in, without it.")
            return io.NodeOutput(patched, positive, f"not applied: {why}")

        tokenizer = _ask("text_tokenizer", clip)
        if tokenizer is None:
            return refuse("this text encoder is not one a model module can place words in")
        clean = markup["clean"]
        offsets = list(tokenizer(clean, add_special_tokens=False,
                                 return_offsets_mapping=True)["offset_mapping"])
        cond, meta = positive[0][0], positive[0][1]
        cond_len = int(cond.shape[1])
        rows = _ask("prompt_rows", meta, cond_len, len(offsets))
        if rows is None:
            return refuse("the encoded prompt does not match this text -- was the clean "
                          "text from Prompt Markup what got encoded?")
        said = []
        if markup.get("extra"):
            _say(f"[a|b] takes {mk.MAX_OPTIONS} options; the {markup['extra']} after them are ignored")
        if variants:
            got = _variant_rows(rows, clip, tokenizer, variants)
            if got is None:
                _say("a [a|b] option's prompt could not be placed; the options are off")
            else:
                whole = []
                for rows_k in got:                     # the full conditioning, as encoded, with this option's prompt rows
                    raw = cond.clone()
                    raw[:, rows[0]:rows[1]] = rows_k.to(device=raw.device, dtype=raw.dtype)
                    whole.append(raw)
                dit_hooks.publish(patched, SIMUL, {"options": whole, "cond_len": cond_len})
                said.append(f"{len(got)} simultaneous option(s)")
        w_spans = mk.token_spans(offsets, weighted or [])
        if w_spans:
            cls._install(patched, w_spans, rows[0], cond_len)
            # Where the phrases sit, for anything that wants to look at them (the phrase probe).
            dit_hooks.publish(patched, PHRASES, {
                "base": rows[0], "cond_len": cond_len, "spans": [(a, b) for a, b, *_ in w_spans]})
            said.append(f"{len(w_spans)} weight(s)")
        return io.NodeOutput(patched, positive, ", ".join(said) or "nothing placeable")

    @staticmethod
    def _install(patcher, w_spans, base, cond_len):
        cache = {}

        def override(func, q, k, v, heads, *args, mask=None, **kwargs):
            def run(m):
                return func(q, k, v, heads, *args, mask=m, **kwargs)

            # Only the packed self-attention: the prompt plus the streams after it.
            # A sequence exactly as long as the prompt is the text-only token
            # refiner, which must never see the bias.
            if k.ndim != 4 or q.shape[2] != k.shape[2] or k.shape[2] <= cond_len:
                return run(mask)
            try:
                key = (int(k.shape[2]), k.device, k.dtype)
                if key not in cache:
                    cache[key] = mk.key_bias(w_spans, base, *key)
                bias = cache[key]
                return run(mask if bias is None else (bias if mask is None else mask + bias))
            except Exception as exc:             # noqa: BLE001 -- markup never costs the step
                _say(f"markup failed inside attention and was dropped ({exc})")
                return run(mask)

        dit_hooks.add_attention_override(patcher, KEY, override)