Upload ComfyUI-ShotSplitter/tools/ui_to_api.py with huggingface_hub

ComfyUI-ShotSplitter/tools/ui_to_api.py

@@ -0,0 +1,398 @@
+"""Convert a ComfyUI UI workflow JSON to the API prompt format.
+
+Best-effort: handles bypassed nodes (mode=4), frontend-only Note/Label/Reroute,
+INT seed control_after_generate hidden widgets, and basic link-to-widget mapping.
+Not a perfect replica of the frontend's graphToPrompt, but enough to run a real
+workflow via POST /prompt for E2E smoke testing.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+import urllib.request
+from typing import Any, Dict, List, Optional, Tuple
+
+FRONTEND_ONLY_TYPES = {
+    "Note",
+    "MarkdownNote",
+    "Label (rgthree)",
+    "Reroute",
+    "PrimitiveNode",  # rarely — frontend injects value into consumer
+    "Fast Groups Bypasser (rgthree)",
+    "Fast Groups Muter (rgthree)",
+    "Image Comparer (rgthree)",
+    "ozW_BrandingNode",
+    "PreviewAny",
+    # KJNodes virtual variables — resolved via set_by_name in resolve_link.
+    "SetNode",
+    "GetNode",
+    # rgthree anywhere node — frontend-only dispatcher
+    "Anything Everywhere",
+    "mxSlider",
+}
+
+# Nodes with a custom frontend widget layout where widgets_values position does
+# NOT cleanly map to INPUT_TYPES declaration order. For those, OMIT all widget
+# inputs in the API call and let the node fall back to its declared defaults.
+IGNORE_WIDGETS_TYPES = {
+    "ozW_SynthesizeAuthenticMetadata",
+}
+
+# Nodes whose widgets_values contain a dynamic list of LoRA entries (each a
+# dict with on/lora/strength). The frontend serializes these as `lora_1`,
+# `lora_2`, ... inputs in the API prompt via FlexibleOptionalInputType. Our
+# converter must replicate that or the LoRAs are silently dropped — the node
+# runs with no loras and produces a bypass identity model patch (i.e. the
+# LoRA appears to not work at all).
+DYNAMIC_LORA_NODE_TYPES = {
+    "Power Lora Loader (rgthree)",
+}
+
+# Widget entries that are UI-only (videopreview state, previews, etc.) and
+# should never be sent to the API — they are dict values with specific keys
+# that identify them as frontend state.
+def _is_ui_widget_state(w):
+    if not isinstance(w, dict):
+        return False
+    ui_keys = ("hidden", "paused", "params", "videopreview")
+    return any(k in w for k in ui_keys)
+
+# Primitive widget types — these ARE allowed to consume widgets_values entries.
+# Anything else (IMAGE, MODEL, CLIP, CONDITIONING, LATENT, MASK, VAE, SEGS, ...)
+# is a connection-only input — never consumes from widgets_values.
+PRIMITIVE_TYPES = {"INT", "FLOAT", "STRING", "BOOLEAN", "COMBO"}
+
+
+def fetch_object_info(server: str = "http://localhost:8188") -> Dict[str, Any]:
+    with urllib.request.urlopen(f"{server}/object_info", timeout=30) as r:
+        return json.load(r)
+
+
+SERVER_INJECTED_HIDDEN_TYPES = {
+    "UNIQUE_ID",
+    "EXTRA_PNGINFO",
+    "PROMPT",
+    "DYNPROMPT",
+    "AUTH_TOKEN_COMFY_ORG",
+    "API_KEY_COMFY_ORG",
+}
+
+
+def ordered_input_names(spec: Dict[str, Any]) -> List[Tuple[str, Any]]:
+    """Return [(input_name, input_config), ...] in declaration order.
+
+    Includes `hidden` inputs whose type is NOT server-injected (e.g.,
+    `prompt_batch_data` of type STRING that Oz_RealityPromptGenerator uses).
+    These are stored as widget values in the UI workflow and must be forwarded
+    via the API prompt; otherwise they fall back to defaults at execution time.
+    """
+    inputs = spec.get("input", {}) or {}
+    result: List[Tuple[str, Any]] = []
+    for section in ("required", "optional"):
+        for name, cfg in (inputs.get(section) or {}).items():
+            result.append((name, cfg))
+    for name, cfg in (inputs.get("hidden") or {}).items():
+        # Skip inputs that ComfyUI's server injects automatically.
+        if isinstance(cfg, str) and cfg in SERVER_INJECTED_HIDDEN_TYPES:
+            continue
+        if isinstance(cfg, (list, tuple)) and cfg and cfg[0] in SERVER_INJECTED_HIDDEN_TYPES:
+            continue
+        result.append((name, cfg))
+    return result
+
+
+def _first_type(cfg: Any) -> str:
+    """Extract the type string from an INPUT_TYPES value."""
+    if isinstance(cfg, list) and cfg:
+        t = cfg[0]
+        if isinstance(t, str):
+            return t
+        if isinstance(t, list):
+            return "COMBO"
+    if isinstance(cfg, tuple) and cfg:
+        return str(cfg[0])
+    return "UNKNOWN"
+
+
+def convert(
+    ui_wf: Dict[str, Any],
+    object_info: Dict[str, Any],
+    verbose: bool = False,
+) -> Dict[str, Any]:
+    api: Dict[str, Any] = {}
+    links_by_id: Dict[int, List[Any]] = {L[0]: L for L in ui_wf.get("links", [])}
+    nodes_by_id: Dict[int, Dict[str, Any]] = {n["id"]: n for n in ui_wf["nodes"]}
+
+    # Build SetNode name → source link mapping for GetNode resolution.
+    # KJNodes / cg-use-everywhere SetNode stashes its input under a name;
+    # GetNode of the same name reads it back. These are virtual on the server —
+    # the frontend rewrites connections to bypass them entirely. We do the
+    # same here.
+    set_by_name: Dict[str, int] = {}  # name -> SetNode's incoming link id
+    for n in ui_wf.get("nodes", []):
+        if n.get("type") != "SetNode":
+            continue
+        wv = n.get("widgets_values") or []
+        name = wv[0] if wv and isinstance(wv[0], str) else None
+        if not name:
+            continue
+        for inp in n.get("inputs") or []:
+            if inp.get("link") is not None:
+                set_by_name[name] = inp["link"]
+                break
+
+    # Build broadcast map from all "Anything Everywhere" nodes. These are
+    # rgthree/cg-use-everywhere pseudo-nodes that broadcast a value of a given
+    # type to every consumer of that type in the graph. The frontend resolves
+    # these at queue time; we must do the same for headless conversion.
+    broadcast_by_type: Dict[str, Tuple[str, int]] = {}
+    for n in ui_wf.get("nodes", []):
+        if n.get("type") not in ("Anything Everywhere", "Anything Everywhere?",
+                                  "Anything Everywhere3", "Seed Everywhere",
+                                  "Prompts Everywhere"):
+            continue
+        for inp in n.get("inputs") or []:
+            lid = inp.get("link")
+            if lid is None:
+                continue
+            itype = inp.get("type")
+            if not itype or itype == "*":
+                continue
+            row = links_by_id.get(lid)
+            if not row:
+                continue
+            src_id, src_slot = row[1], row[2]
+            # Walk through Reroute/GetNode/SetNode like resolve_link does
+            # (simple version — just take the direct source for now)
+            broadcast_by_type[itype] = (str(src_id), src_slot)
+
+    skipped_bypassed: List[int] = []
+    skipped_frontend: List[int] = []
+    skipped_unknown: List[int] = []
+
+    def resolve_link(
+        link_id: int, visited: Optional[set] = None
+    ) -> Optional[Tuple[str, int]]:
+        """Walk through Reroute and bypassed (mode=4) nodes to find a real source.
+
+        For a bypassed node, we pass through its input of the SAME type at the
+        same slot index — rgthree-style bypass semantics.
+        """
+        if visited is None:
+            visited = set()
+        if link_id in visited:
+            return None
+        visited.add(link_id)
+        row = links_by_id.get(link_id)
+        if not row:
+            return None
+        _lid, src_id, src_slot, _tgt_id, _tgt_slot, src_type = row
+        src_node = nodes_by_id.get(src_id)
+        if not src_node:
+            return None
+
+        # Reroute: one input, one output, walk upstream
+        if src_node.get("type") == "Reroute":
+            src_inputs = src_node.get("inputs") or []
+            if src_inputs and src_inputs[0].get("link") is not None:
+                return resolve_link(src_inputs[0]["link"], visited)
+            return None
+
+        # GetNode: virtual KJNodes variable — look up the matching SetNode
+        # by name and recursively resolve that SetNode's stored source link.
+        if src_node.get("type") == "GetNode":
+            wv = src_node.get("widgets_values") or []
+            name = wv[0] if wv and isinstance(wv[0], str) else None
+            if name and name in set_by_name:
+                return resolve_link(set_by_name[name], visited)
+            return None
+
+        # SetNode: pass-through to whatever is stored (its first input link).
+        # Some SetNodes act as both stash AND passthrough — downstream edges
+        # can go directly from SetNode's output.
+        if src_node.get("type") == "SetNode":
+            for inp in src_node.get("inputs") or []:
+                if inp.get("link") is not None:
+                    return resolve_link(inp["link"], visited)
+            return None
+
+        # Bypassed node: pass-through via same-type input (prefer same slot index)
+        if src_node.get("mode") == 4:
+            src_inputs = src_node.get("inputs") or []
+            candidates = [
+                (i, inp)
+                for i, inp in enumerate(src_inputs)
+                if inp.get("type") == src_type and inp.get("link") is not None
+            ]
+            if not candidates:
+                return None
+            # Prefer the input at the same index as the output slot, else first
+            match = next(
+                (inp for i, inp in candidates if i == src_slot),
+                candidates[0][1],
+            )
+            return resolve_link(match["link"], visited)
+
+        return (str(src_id), src_slot)
+
+    for node in ui_wf["nodes"]:
+        ntype = node.get("type", "")
+        nid = str(node["id"])
+
+        if node.get("mode") == 4:
+            skipped_bypassed.append(node["id"])
+            continue
+        if ntype in FRONTEND_ONLY_TYPES:
+            skipped_frontend.append(node["id"])
+            continue
+
+        spec = object_info.get(ntype)
+        if not spec:
+            skipped_unknown.append(node["id"])
+            if verbose:
+                print(f"WARN: unknown type {ntype} id={nid}")
+            continue
+
+        ordered = ordered_input_names(spec)
+
+        # For nodes with broken widget→inputs mapping, only forward connection
+        # inputs and let the node use declared defaults for every widget.
+        ignore_widgets = ntype in IGNORE_WIDGETS_TYPES
+
+        # Fresh per-node dicts — MUST be reset here (not below) because the
+        # dict-typed widgets_values branch also consumes api_inputs and would
+        # otherwise inherit stale state from the previous iteration.
+        api_inputs: Dict[str, Any] = {}
+        connection_by_name: Dict[str, int] = {}
+        for inp in node.get("inputs") or []:
+            if inp.get("link") is not None:
+                connection_by_name[inp.get("name")] = inp["link"]
+
+        # Dict-typed widgets_values (new ComfyUI schema) — map by NAME directly
+        # into api_inputs, skipping UI-only state entries like videopreview.
+        # Used by VHS_LoadVideo and some other nodes that moved to dict form.
+        wv_raw = node.get("widgets_values")
+        if isinstance(wv_raw, dict):
+            if not ignore_widgets:
+                for wk, wv in wv_raw.items():
+                    if _is_ui_widget_state(wv):
+                        continue
+                    # Only forward names that the node declares as inputs
+                    if any(name == wk for name, _ in ordered):
+                        # Skip if the input is already a connection
+                        if wk not in {n_.get("name") for n_ in (node.get("inputs") or []) if n_.get("link") is not None}:
+                            api_inputs[wk] = wv
+            # Still process connections + broadcast for unconnected inputs
+            for name, cfg in ordered:
+                if name in connection_by_name:
+                    resolved = resolve_link(connection_by_name[name])
+                    if resolved is not None:
+                        api_inputs[name] = list(resolved)
+                elif name not in api_inputs:
+                    type_str = _first_type(cfg)
+                    if (type_str not in PRIMITIVE_TYPES
+                            and type_str in broadcast_by_type):
+                        api_inputs[name] = list(broadcast_by_type[type_str])
+            if ntype in DYNAMIC_LORA_NODE_TYPES:
+                lora_idx = 1
+                for w in wv_raw.values():
+                    if (isinstance(w, dict) and "on" in w and "lora" in w
+                            and "strength" in w):
+                        api_inputs[f"lora_{lora_idx}"] = {
+                            "on": w.get("on", True),
+                            "lora": w.get("lora", ""),
+                            "strength": w.get("strength", 1),
+                            "strengthTwo": w.get("strengthTwo"),
+                        }
+                        lora_idx += 1
+            api[nid] = {"class_type": ntype, "inputs": api_inputs}
+            continue
+
+        # connection_by_name already built above before dict branch
+
+        wv = list(node.get("widgets_values") or [])
+        wi = 0  # widget index
+
+        # Drop UI-only state entries (videopreview dict) from the positional
+        # list so we don't accidentally map them onto a real input slot.
+        wv = [w for w in wv if not _is_ui_widget_state(w)]
+
+        for name, cfg in ordered:
+            type_str = _first_type(cfg)
+            is_primitive = type_str in PRIMITIVE_TYPES
+
+            connected = name in connection_by_name
+
+            if connected:
+                resolved = resolve_link(connection_by_name[name])
+                if resolved is not None:
+                    api_inputs[name] = list(resolved)
+                # else: missing upstream — omit
+            elif not is_primitive and type_str in broadcast_by_type:
+                # Anything Everywhere broadcast: inject the shared source for
+                # this type if the input is unconnected.
+                api_inputs[name] = list(broadcast_by_type[type_str])
+
+            # Primitive inputs have a SLOT in widgets_values even when converted
+            # to a connection (ComfyUI keeps the stale value). So ALWAYS advance
+            # the widget pointer for primitive inputs.
+            if is_primitive and wi < len(wv):
+                if not connected and not ignore_widgets:
+                    api_inputs[name] = wv[wi]
+                wi += 1
+                # Handle control_after_generate hidden widget for seeds
+                opts = cfg[1] if isinstance(cfg, list) and len(cfg) > 1 else {}
+                if isinstance(opts, dict) and (
+                    opts.get("control_after_generate") or name == "seed"
+                ):
+                    if wi < len(wv) and isinstance(wv[wi], str) and wv[wi] in (
+                        "fixed", "increment", "decrement", "randomize"
+                    ):
+                        wi += 1
+
+        # Ignore hidden inputs (unique_id, extra_pnginfo) — server injects them
+
+        # Dynamic LoRA entries for nodes like rgthree Power Lora Loader.
+        # These are stored as dict widgets at arbitrary positions in
+        # widgets_values; extract each `{on, lora, strength}` and expose as
+        # lora_1, lora_2, ... so the node sees them in its **kwargs.
+        if ntype in DYNAMIC_LORA_NODE_TYPES:
+            lora_idx = 1
+            for w in node.get("widgets_values") or []:
+                if (isinstance(w, dict) and "on" in w and "lora" in w
+                        and "strength" in w):
+                    api_inputs[f"lora_{lora_idx}"] = {
+                        "on": w.get("on", True),
+                        "lora": w.get("lora", ""),
+                        "strength": w.get("strength", 1),
+                        "strengthTwo": w.get("strengthTwo"),
+                    }
+                    lora_idx += 1
+
+        api[nid] = {"class_type": ntype, "inputs": api_inputs}
+
+    if verbose:
+        print(f"bypassed skipped: {len(skipped_bypassed)}")
+        print(f"frontend skipped: {len(skipped_frontend)}")
+        print(f"unknown skipped: {len(skipped_unknown)}")
+        print(f"api nodes: {len(api)}")
+    return api
+
+
+def main() -> None:
+    if len(sys.argv) < 2:
+        print("usage: ui_to_api.py <workflow.json> [output.json]")
+        sys.exit(2)
+    with open(sys.argv[1], encoding="utf-8") as f:
+        ui = json.load(f)
+    info = fetch_object_info()
+    api = convert(ui, info, verbose=True)
+    out = sys.argv[2] if len(sys.argv) > 2 else "api_prompt.json"
+    with open(out, "w", encoding="utf-8") as f:
+        json.dump(api, f, ensure_ascii=False, indent=2)
+    print(f"wrote {out}: {len(api)} nodes")
+
+
+if __name__ == "__main__":
+    main()