olivv-cs-ComfyUI-FunPack/modules/sampling/phrase_probe/__init__.py

"""Phrase probe: where does the model read a bracketed phrase, and where does it matter?

MEASUREMENT ONLY -- nothing about a generation changes. Switched on from Settings >
Refinement & Taste (off by default: it is research data, and the response map costs one
extra model forward per phrase per call).

The phrases are the ones Prompt Markup placed -- weighted `(word:1.5)` and timed
`[phrase@2-4]` -- which it publishes into the model's transformer_options. That is read
at run time, because this installs BEFORE the markup node runs in the pipeline.

v4's first result (blocks 15-20 read the text, 29-34 and block 43 commit it, blocks 0-6 are
text-blind; two different shortcuts gave the same map) came from a run with steering on.
Validated on a rental in v5.
"""

import torch

from ..._core import dit_hooks, log, streams
from . import measure

ID = "phrase_probe"
TITLE = "Phrase probe"
MOUNT = "generation.sampling"
STAGE = "post"
CATEGORY = "system"
STATUS = "experimental"
REQUIRES = ["dit_block_hooks"]
# Innermost on the model, so every call it sees is a real forward: the extra forwards of
# late-branch guidance and the taste slider (flagged as probes) pass through it untouched
# instead of being measured as the picture.
AFTER = ["late_branch", "score_slider", "sharpen"]


def _say(message):
    log.once(f"{ID}:{message}", log.ALERT, "FunUI Phrase probe", message)


def install(patcher, values, key):
    if not measure.enabled():
        measure.problem = None
        return None                              # off: no wrappers, no cost, nothing to disturb
    n = dit_hooks.block_count(patcher)
    if n == 0:
        measure.problem = "this model's blocks can't be read, so nothing is measured"
        if measure.enabled():
            _say("recording is on but " + measure.problem)
        return None
    measure.problem = None
    live = {"on": False, "state": None}

    def hook_for(block):
        def hook(args, extra):
            st = live["state"]
            if not live["on"] or st is None:
                return extra["original_block"](args)
            return measure.block_hook(st, block)(args, extra)
        return hook

    for b in range(n):
        dit_hooks.add_block_hook(patcher, key, b, hook_for(b))

    def attention(func, *a, **kw):
        st = live["state"]
        if not live["on"] or st is None:
            return func(*a, **kw)
        return measure.attention_override(st)(func, *a, **kw)

    dit_hooks.add_attention_override(patcher, key, attention)
    # Masked forwards allocate differently from the visible one: see dit_hooks.without_compiler.
    dit_hooks.without_compiler(patcher, key)

    from comfy.patcher_extension import WrappersMP

    def outer(executor, *args, **kwargs):
        live["on"] = measure.enabled()
        live["state"] = None
        try:
            return executor(*args, **kwargs)
        finally:
            st = live["state"]
            if live["on"] and st is not None:
                res = measure.result(st)
                measure.save(res)
                for p in measure.peaks(res, top=3):
                    rt = ", ".join(f"b{b} {v * 100:.1f}%" for b, v in p["read_top"])
                    gt = ", ".join(f"b{b} +{v:.3f}" for b, v in p["response_growth_top"])
                    log.once(f"{ID}:run:{p['phrase']}", log.INFO, "FunUI Phrase probe",
                             f"Active | phrase {p['phrase']} ({p['tokens']} tokens): read most at "
                             f"{rt or 'n/a'}; response grows most at {gt or 'n/a'}")
            elif live["on"]:
                measure.problem = ("recording is on but this run had no phrases to measure: put a "
                                   "(word:1.5) or [phrase@2-4] in the prompt")
                log.once(f"{ID}:none", log.INFO, "FunUI Phrase probe", "Inactive | " + measure.problem)
            else:
                measure.problem = None

    patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, outer)

    def apply_model(executor, x, t, *args, **kwargs):
        named = streams.model_args(args, kwargs)
        to = named.get("transformer_options") or {}
        if not live["on"] or dit_hooks.probing(to):
            return executor(x, t, *args, **kwargs)
        st = live["state"]
        if st is None:
            phrases = measure.phrases_of(to.get("funui_markup_phrases"))
            if not phrases:
                return executor(x, t, *args, **kwargs)
            st = live["state"] = measure.State(phrases, to["funui_markup_phrases"]["cond_len"])
            log.once(f"{ID}:cost:{len(phrases)}", log.INFO, "FunUI Phrase probe",
                     f"{len(phrases)} phrase(s): the reading map is free, the response map costs "
                     f"{len(phrases)} extra forward(s) per model call this run")
        st.mode, st.block = -1, 0
        try:
            out = executor(x, t, *args, **kwargs)
        finally:
            st.mode = None
        st.calls += 1
        probe_args, probe_kwargs = streams.with_options(args, kwargs, {**to, dit_hooks.PROBE: True})
        for i in range(len(st.phrases)):
            st.mode, st.block = i, 0
            try:
                executor(x, t, *probe_args, **probe_kwargs)
            except Exception as exc:                  # noqa: BLE001 -- the probe's own pass never costs the step
                _say(f"a masked pass failed ({type(exc).__name__}: {exc}); this phrase's response is not measured")
            finally:
                st.mode = None
        st.visible = {}
        return out

    patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model)
    return "measures where each marked phrase is read and where it matters (off unless recording is switched on)"


def routes(table, base, web):
    @table.get(base + "/status")
    async def _status(_req):
        return web.json_response({"enabled": measure.enabled(), "problem": measure.problem,
                                  "latest": measure.latest(), "peaks": measure.peaks(measure.latest())})

    @table.post(base + "/enabled")
    async def _enabled(req):
        try:
            body = await req.json()
        except Exception:                             # noqa: BLE001
            body = {}
        on = measure.set_enabled((body or {}).get("enabled") is True)
        return web.json_response({"enabled": on, "problem": measure.problem, "latest": measure.latest(),
                                  "peaks": measure.peaks(measure.latest())})

    @table.post(base + "/clear")
    async def _clear(_req):
        measure.clear()
        return web.json_response({"enabled": measure.enabled(), "problem": measure.problem, "latest": None,
                                  "peaks": []})

    @table.get(base + "/export")
    async def _export(_req):
        res = measure.latest()
        if res is None:
            return web.json_response({"why": "nothing recorded yet"}, status=404)
        return web.json_response(res, headers={"Content-Disposition":
                                               'attachment; filename="phrase_probe.json"'})


def cache_key():
    # The node that installs this is cached by ComfyUI: the switch must be part of what re-runs it.
    return "on" if measure.enabled() else "off"


PROVIDES = {"modifier": install, "routes": routes, "cache_key": cache_key}