DemonGatanjieu-Anomalous_Model_Browser/api/mcp_fill.py

"""run_workflow's blanks: a daily run in one call, so a small model need not read and edit the canvas.

With `workflow` (a saved workflow's name, or one of its 方案 as "workflow / plan": a parameter set saved
under it, applied as its Apply button does, or a gallery picture, whose own workflow is opened: `opening`)
the run opens it in the task's tab the first time and reuses that tab after (the page marks a tab
open_workflow opened with the workflow's file, or "picture:" and the picture); what
the AI fills in goes into it first (web/modules/mcp_canvas_edit.js fill_slots): the positive and
negative prompt boxes, the seed of every sampler, the empty latent's width and height, the
picture its LoadImage reads (from the input or the output folder, found here, never copied;
image_node says which LoadImage when there are several), and `values`: any other setting by its
name, a third-party node's as well, for the one node that has it ({node id: {...}} when several
do; set_values takes the same). No fixed list of blanks grows here. The rest stays as saved. Without
`workflow` the blanks go into the canvas the run would run. All of it is checked before anything
changes, and it is one Ctrl+Z step logged as the AI's.
"""

import os

from . import ai_task_canvas, ai_tasks
from .ai_pictures import to_input
from .mcp_bridge import PageRefused, PageUnavailable, ask_page
from .mcp_tools import IMAGE_TYPES, ToolError, _resolve_folder_image, _resolve_image, plans
from .recipe_images import _embedded_workflow_payload

SLOTS = ("prompt", "negative", "seed", "width", "height", "image", "image_node", "values")
MAX_VALUES = 60
MAX_TEXT = 20000


def flat_values(values):
    """{setting name: value}, checked for shape and size; the page checks what each setting takes."""
    if values is None:
        return {}
    if not isinstance(values, dict) or len(values) > MAX_VALUES:
        raise ToolError(f"values is an object {{setting name: value}} of up to {MAX_VALUES} settings.")
    for name, value in values.items():
        if len(name) > 100 or not (isinstance(value, (bool, int, float)) or (isinstance(value, str) and len(value) <= MAX_TEXT)):
            raise ToolError(f"The value of {name[:100]} is a number, true/false or text of up to {MAX_TEXT} characters.")
    return values


def node_values(values):
    """Settings by name ({setting name: value}) or for one node ({node id: {setting name: value}})."""
    if not isinstance(values, dict) or not values:
        raise ToolError('values is {setting name: value}, or {"<node id>": {setting name: value}} for one node.')
    nested = {key: value for key, value in values.items() if isinstance(value, dict)}
    for key, settings in nested.items():
        if not str(key).isdigit():
            raise ToolError(f"values.{str(key)[:40]}: only a node id takes an object of settings.")
        flat_values(settings)
    flat_values({key: value for key, value in values.items() if key not in nested})
    if len(values) - len(nested) + sum(len(settings) for settings in nested.values()) > MAX_VALUES:
        raise ToolError(f"values sets up to {MAX_VALUES} settings at once.")
    return values


def saved_workflow(name):
    """(file name, name, plan) of the user's saved workflow called `name`, or of its 方案 "workflow / plan"
    (plan: the parameter set's file name, else None), or of the example called that."""
    from . import example_workflows
    from .recipe_store import _list_recipes, get_recipes_dir

    name = str(name or "").strip()
    files = {}
    for recipe in _list_recipes(get_recipes_dir()):
        files.setdefault(recipe["data"].get("name") or recipe["filename"], recipe["filename"])
    if name and name in files:
        return files[name], name, None
    example = example_workflows.file_of(name)
    if example:
        return example, name, None
    workflow, _, plan = name.rpartition(" / ")
    found = next((file for file, plan_name in plans().get(files.get(workflow), []) if plan_name == plan), None)
    if not found:
        raise ToolError(f"No saved workflow or plan called {name}; list_workflows names them.")
    return files[workflow], name, found


def picture_workflow(name):
    """{subfolder, filename} of gallery picture `name` (in the output folder) when the name is a
    picture's and it carries a workflow; None when the name is no picture's."""
    if os.path.splitext(str(name or "").strip())[1].lower() not in IMAGE_TYPES:
        return None
    rel, path = _resolve_image(str(name).strip())
    if not (_embedded_workflow_payload(path) or {}).get("workflow"):
        raise ToolError(f"{rel} carries no workflow to open: only the pictures ComfyUI saves with one do.")
    subfolder, _, filename = rel.rpartition("/")
    return {"subfolder": subfolder, "filename": filename}


def opening(name):
    """(what the page's open_workflow takes, the mark it leaves on the tab) for `name`: a saved workflow,
    one of its plans or an example (saved_workflow), else a gallery picture's own workflow."""
    try:
        filename, shown, plan = saved_workflow(name)
    except ToolError:
        picture = picture_workflow(name)  # not a saved one: a picture's own, to change what made it
        if not picture:
            raise
        rel = "/".join(part for part in (picture["subfolder"], picture["filename"]) if part)
        return {"picture": picture, "name": os.path.splitext(picture["filename"])[0]}, f"picture:{rel}"
    return {"filename": filename, "name": shown, "plan": plan}, (f"{filename}#{plan}" if plan else filename)


def _picture(image):
    """The value a LoadImage takes for picture `image`: in the input folder as it is, from the
    output folder its copy there (ai_pictures.to_input)."""
    for folder in ("input", "output"):
        try:
            rel, _path = _resolve_folder_image(image, folder)
        except ToolError:
            continue
        return rel if folder == "input" else to_input(rel)
    raise ToolError(f"{image} is in neither ComfyUI's input nor its output folder; search_images lists the pictures in the output.")


def _page(action, args, timeout=30):
    try:
        return ask_page(action, args, timeout)
    except (PageUnavailable, PageRefused) as error:
        raise ToolError(str(error))


def fill(ctx, workflow=None, **slots):
    """Opens `workflow` for the task (once) and fills `slots` in; what was set, or None when nothing was asked."""
    slots = {key: value for key, value in slots.items() if key in SLOTS and value not in (None, "")}
    if not workflow and not slots:
        return None
    if "image" in slots:
        slots["image"] = _picture(slots["image"])
    if "values" in slots:
        node_values(slots["values"])
    args, opened = opening(workflow) if workflow else ({}, None)  # opened: what the page marks the tab with
    name = args.get("name")
    task = ai_tasks.get_task(ctx.get("task")) if ctx.get("task") else None
    no_tab_yet = bool(task) and ai_task_canvas.mode_of(task) == "new" and not (task.get("canvas") or {}).get("tab")
    filled = None if workflow and no_tab_yet else _page("fill_slots", {"filename": opened, **slots})
    if workflow and (filled is None or filled.get("needs_open")):
        _page("open_workflow", args, timeout=60)
        filled = _page("fill_slots", {"filename": opened, **slots})
        if filled.get("needs_open"):
            raise ToolError(f"The saved workflow {name} did not open; nothing was run.")
    return {**({"workflow": name} if name else {}), **(filled.get("set") or {})}