DemonGatanjieu-Anomalous_Model_Browser/api/ai_run_map.py

"""The shape of a workflow run in a few boxes, so the user sees what the AI really ran: read from
the prompt ComfyUI queued (ai_task_runs.py keeps it on the run's step as "map"), never from
what the AI says. Plumbing (reroutes, primitives, notes) is left out and wired through; each
box has a role (model, lora, prompt, size, sampler, image, upscale, decode, save, other) and a
few facts the page words in the user's language (ui_ai_run_map.js); the links carry ComfyUI's
types. changes() says what differs from the run before it: a setting, a node added or gone.
"""

import os

from .model_constants import MODEL_EXTENSIONS

MAX_BOXES = 24  # what a stage diagram holds (ai_panels.MAX_NODES)
MAX_EDGES = 40
MAX_CHANGES = 10
HIDDEN = {"Reroute", "Note", "MarkdownNote", "PrimitiveNode", "PrimitiveInt", "PrimitiveFloat", "PrimitiveString",
          "PrimitiveStringMultiline", "PrimitiveBoolean"}
SEEDS = {"seed", "noise_seed"}
SAMPLER_FACTS = ("steps", "cfg", "sampler_name", "scheduler", "denoise")
TEXT_CHARS = 80
VALUE_CHARS = 60


def _linked(value):
    return isinstance(value, list) and len(value) == 2 and isinstance(value[0], (str, int))


def _literal(value):
    return isinstance(value, (str, int, float, bool)) and not isinstance(value, list)


def _file_name(value):
    return os.path.splitext(os.path.basename(str(value).replace("\\", "/")))[0]


def _output_type(class_type, index):
    """ComfyUI's type of output `index` of `class_type` ("MODEL", "LATENT"...), or ""."""
    try:
        import nodes
        types = getattr(nodes.NODE_CLASS_MAPPINGS.get(class_type), "RETURN_TYPES", None) or ()
        return str(types[int(index)]) if int(index) < len(types) else ""
    except Exception:  # noqa: BLE001 - no ComfyUI (tests) or an odd node: the link stays unnamed
        return ""


def _role(class_type, inputs):
    name = class_type.lower()
    if "lora_name" in inputs:
        return "lora"
    if _linked(inputs.get("latent_image")) or "sampler" in name and not name.endswith("select"):
        return "sampler"
    if name == "loadimage" or name.startswith("loadimage"):
        return "image"
    if "upscale" in name:
        return "upscale"
    if "loader" in name:
        return "model"
    if isinstance(inputs.get("text"), str) or "textencode" in name:
        return "prompt"
    if name.startswith("empty") and "latent" in name:
        return "size"
    if name.startswith("vaedecode"):
        return "decode"
    if name.startswith(("save", "preview")):
        return "save"
    return "other"


def _facts(role, inputs):
    literal = {key: value for key, value in inputs.items() if _literal(value)}
    facts = {}
    model = next((value for value in literal.values() if isinstance(value, str) and value.lower().endswith(MODEL_EXTENSIONS)), None)
    if model:
        facts["file"] = _file_name(model)[:VALUE_CHARS]
    if role == "lora":
        for key in ("strength_model", "strength"):
            if isinstance(literal.get(key), (int, float)):
                facts["strength"] = round(float(literal[key]), 3)
                break
    elif role == "prompt":
        text = next((value for key, value in literal.items() if isinstance(value, str) and key.startswith("text")), "")
        facts["text"] = " ".join(text.split())[:TEXT_CHARS]
    elif role == "sampler":
        facts.update({key: literal[key] for key in SAMPLER_FACTS if key in literal})
    elif role == "image":
        facts["file"] = str(literal.get("image") or "")[:VALUE_CHARS]
    elif role == "save":
        facts["prefix"] = str(literal.get("filename_prefix") or "")[:VALUE_CHARS]
    if all(isinstance(literal.get(key), (int, float)) and literal[key] > 0 for key in ("width", "height")):
        facts.update(width=int(literal["width"]), height=int(literal["height"]))
        if isinstance(literal.get("batch_size"), int) and literal["batch_size"] > 1:
            facts["batch"] = literal["batch_size"]
    if role == "other" and not facts:
        facts["values"] = [[key, str(value)[:VALUE_CHARS]] for key, value in literal.items()
                           if key not in SEEDS and str(value).strip()][:2]
    return facts


def map_of(prompt):
    """{"boxes": [{id, type, title, role, facts}], "links": [{from, to, type}], "params": {...}} of API prompt `prompt`."""
    nodes = {str(key): node for key, node in (prompt or {}).items() if isinstance(node, dict) and node.get("class_type")}
    shown = [key for key in nodes if nodes[key]["class_type"] not in HIDDEN]
    if len(shown) > MAX_BOXES:  # keep what makes the picture; drop the helpers first
        rank = {"other": 2, "save": 1}
        shown = sorted(shown, key=lambda key: rank.get(_role(nodes[key]["class_type"], nodes[key].get("inputs") or {}), 0))[:MAX_BOXES]
    keep = set(shown)

    def source(link, seen=()):
        """The shown node (and the output's type) a link comes from, wired through hidden nodes."""
        key, index = str(link[0]), link[1]
        if key in keep:
            return key, _output_type(nodes[key]["class_type"], index)
        node = nodes.get(key)
        if not node or key in seen:
            return None, ""
        upstream = next((value for value in (node.get("inputs") or {}).values() if _linked(value)), None)
        return source(upstream, (*seen, key)) if upstream else (None, "")

    boxes, links, params = [], [], {}
    for key in sorted(shown, key=lambda item: (len(item), item)):
        node = nodes[key]
        inputs = node.get("inputs") or {}
        role = _role(node["class_type"], inputs)
        boxes.append({"id": key, "type": node["class_type"][:80], "role": role, "facts": _facts(role, inputs),
                      "title": str((node.get("_meta") or {}).get("title") or node["class_type"])[:80]})
        for name, value in inputs.items():
            if _linked(value):
                found, kind = source(value)
                if found and found != key and len(links) < MAX_EDGES:
                    links.append({"from": found, "to": key, "type": kind or name.upper(), "input": name})
            elif _literal(value) and name not in SEEDS:
                params[f"{key}.{name}"] = str(value)[:300]
    _polarity(boxes, links)
    return {"boxes": boxes, "links": links, "params": params}


def _polarity(boxes, links):
    """Marks prompt boxes positive or negative by the sampler input they reach."""
    by_id = {box["id"]: box for box in boxes}
    feeds = {}
    for link in links:
        feeds.setdefault(link["from"], []).append(link)
    for box in boxes:
        if box["role"] != "prompt":
            continue
        todo, seen = [box["id"]], set()
        while todo and "polarity" not in box["facts"]:
            here = todo.pop()
            seen.add(here)
            for link in feeds.get(here, []):
                if link["input"] in ("positive", "negative") and by_id[link["to"]]["role"] == "sampler":
                    box["facts"]["polarity"] = link["input"]
                    break
                if link["to"] not in seen:
                    todo.append(link["to"])


def changes(before, after):
    """What differs in map `after` from map `before`: [{node, title, key, from, to}] or a node added / gone."""
    if not before or not after:
        return []
    titles = {box["id"]: box["title"] for box in before.get("boxes", [])}
    titles.update({box["id"]: box["title"] for box in after.get("boxes", [])})
    old_ids = {box["id"]: box["type"] for box in before.get("boxes", [])}
    new_ids = {box["id"]: box["type"] for box in after.get("boxes", [])}
    out = []
    for key, kind in new_ids.items():
        if old_ids.get(key) != kind:
            out.append({"node": key, "title": titles[key], "change": "added"})
    for key, kind in old_ids.items():
        if new_ids.get(key) != kind:
            out.append({"node": key, "title": titles[key], "change": "removed"})
    gone = {item["node"] for item in out}
    old, new = before.get("params", {}), after.get("params", {})
    for name in sorted(set(old) | set(new), key=lambda item: (len(item.split(".")[0]), item)):
        node, key = name.split(".", 1)
        if node not in gone and old.get(name) != new.get(name):
            out.append({"node": node, "title": titles.get(node, node), "key": key,
                        "from": (old.get(name) or "")[:VALUE_CHARS], "to": (new.get(name) or "")[:VALUE_CHARS]})
    return out[:MAX_CHANGES]