workflows/scripts/run_workflow.py

#!/usr/bin/env python3
"""Run an API-format ComfyUI workflow synchronously and wait for it to finish.

Queues the workflow via POST /prompt, polls GET /history/{prompt_id} until the
job completes, then prints one reference name per produced output file to
stdout ("file.png" or "subfolder/file.png") -- feed those straight to
download_files.py. Exits non-zero on validation or execution failure, with
ComfyUI's error details on stderr.

The workflow must be API format (flat {"<id>": {"class_type": ...}} dict).
Export one from the UI with Workflow > Export (API), or convert a saved UI
JSON with convert_workflow.py.

Runtime patching (so workflow files stay generic):
    --set NODE_ID.INPUT=VALUE   override any node input, e.g. --set 76.image=body.jpg
                                (VALUE is parsed as JSON when possible, else kept as string)
    --prefix PREFIX             set filename_prefix on every node that has one,
                                e.g. --prefix "headswap/%year%%month%%day%-%hour%%minute%%second%"
    --randomize-seeds           randomize every seed/noise_seed input (on by default;
                                disable with --keep-seeds for reproducible reruns)

Example:
    python run_workflow.py --url https://POD-7865.proxy.runpod.net \
        ../api/head-swap-flux-klein.json \
        --set 76.image=body.jpg --set 81.image=head.jpg \
        --prefix "headswap/%year%%month%%day%-%hour%%minute%%second%"
"""
import argparse
import json
import pathlib
import random
import sys
import time
import uuid

import requests

SEED_INPUT_NAMES = ("seed", "noise_seed")


def parse_set(spec):
    """'76.image=body.jpg' -> ('76', 'image', 'body.jpg'). Values parse as JSON when possible."""
    try:
        target, raw = spec.split("=", 1)
        node_id, input_name = target.split(".", 1)
    except ValueError:
        raise argparse.ArgumentTypeError(f"--set expects NODE_ID.INPUT=VALUE, got {spec!r}")
    try:
        value = json.loads(raw)
    except json.JSONDecodeError:
        value = raw
    return node_id, input_name, value


def collect_output_files(history_entry):
    """Yield (ref_name, folder_type) for every file any output node produced."""
    for node_output in history_entry.get("outputs", {}).values():
        for value in node_output.values():
            if not isinstance(value, list):
                continue
            for item in value:
                if isinstance(item, dict) and "filename" in item:
                    name = item["filename"]
                    if item.get("subfolder"):
                        name = f"{item['subfolder']}/{name}"
                    yield name, item.get("type", "output")


def run(base_url, workflow, timeout=900, poll_interval=2.0, quiet=False):
    """Queue an API-format workflow and block until done. Returns the history entry."""
    base_url = base_url.rstrip("/")
    r = requests.post(
        f"{base_url}/prompt",
        json={"prompt": workflow, "client_id": str(uuid.uuid4())},
        timeout=60,
    )
    if r.status_code != 200:
        try:
            detail = json.dumps(r.json(), indent=2)
        except ValueError:
            detail = r.text
        raise RuntimeError(f"queueing failed (HTTP {r.status_code}):\n{detail}")
    prompt_id = r.json()["prompt_id"]
    if not quiet:
        print(f"queued prompt {prompt_id}", file=sys.stderr)

    start = time.monotonic()
    while True:
        elapsed = time.monotonic() - start
        if elapsed > timeout:
            raise RuntimeError(f"timed out after {timeout}s waiting for prompt {prompt_id}")
        h = requests.get(f"{base_url}/history/{prompt_id}", timeout=60).json()
        if prompt_id in h:
            entry = h[prompt_id]
            status = entry.get("status", {})
            if status.get("status_str") == "error":
                messages = json.dumps(status.get("messages", []), indent=2)
                raise RuntimeError(f"workflow execution failed:\n{messages}")
            return entry
        if not quiet:
            print(f"  waiting... {int(elapsed)}s", file=sys.stderr)
        time.sleep(poll_interval)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--url", required=True, help="ComfyUI base URL, e.g. https://POD-7865.proxy.runpod.net")
    ap.add_argument("workflow", help="path to an API-format workflow JSON")
    ap.add_argument("--set", dest="sets", type=parse_set, action="append", default=[],
                    metavar="NODE_ID.INPUT=VALUE", help="override a node input (repeatable)")
    ap.add_argument("--prefix", help="set filename_prefix on every node that has one")
    ap.add_argument("--keep-seeds", action="store_true",
                    help="keep the seeds saved in the workflow instead of randomizing")
    ap.add_argument("--timeout", type=float, default=900, help="max seconds to wait (default 900)")
    ap.add_argument("--poll", type=float, default=2.0, help="poll interval in seconds (default 2)")
    args = ap.parse_args()

    workflow = json.loads(pathlib.Path(args.workflow).read_text())
    if "nodes" in workflow and "class_type" not in next(iter(workflow.values()), {}):
        sys.exit("ERROR: this is a UI-format workflow. Convert it first: "
                 "python convert_workflow.py --url <URL> <workflow.json>")

    for node_id, input_name, value in args.sets:
        if node_id not in workflow:
            sys.exit(f"ERROR: --set references unknown node id {node_id!r}")
        workflow[node_id]["inputs"][input_name] = value
        print(f"set {node_id}.{input_name} = {value!r}", file=sys.stderr)

    if args.prefix:
        for node_id, node in workflow.items():
            if "filename_prefix" in node.get("inputs", {}):
                node["inputs"]["filename_prefix"] = args.prefix
                print(f"set {node_id}.filename_prefix = {args.prefix!r}", file=sys.stderr)

    if not args.keep_seeds:
        for node_id, node in workflow.items():
            for name in SEED_INPUT_NAMES:
                if isinstance(node.get("inputs", {}).get(name), int):
                    node["inputs"][name] = random.randrange(2**48)
                    print(f"randomized {node_id}.{name} = {node['inputs'][name]}", file=sys.stderr)

    try:
        entry = run(args.url, workflow, timeout=args.timeout, poll_interval=args.poll)
    except (requests.RequestException, RuntimeError) as e:
        sys.exit(f"ERROR: {e}")

    outputs = list(collect_output_files(entry))
    saved = [name for name, ftype in outputs if ftype == "output"]
    temp = [name for name, ftype in outputs if ftype != "output"]
    if temp:
        print(f"({len(temp)} preview/temp file(s) not listed: {', '.join(temp)})", file=sys.stderr)
    if not saved:
        print("workflow finished but produced no saved output files", file=sys.stderr)
    for name in saved:
        print(name)


if __name__ == "__main__":
    main()