serverless/tests/local_test.py

#!/usr/bin/env python3
"""Drive handler.py directly against a locally running ComfyUI (no RunPod).

    python tests/local_test.py main.jpg [ref.png] [--prompt "..."] [--mode turbo-8]

Saves returned images next to this script and prints timings.
"""
import argparse
import base64
import json
import os
import pathlib
import sys
import time

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
import handler as H  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("images", nargs="+", help="input image file(s): main [ref]")
    ap.add_argument("--prompt", default="replace the outfit of the subject in the first "
                    "image with the outfit in the second image, don't change anything else.")
    ap.add_argument("--mode", default=None, help="only for workflows that support it")
    ap.add_argument("--seed", type=int)
    ap.add_argument("--set", dest="sets", action="append", default=[],
                    metavar="NODE.INPUT=JSONVALUE")
    args = ap.parse_args()

    imgs = []
    for p in args.images:
        imgs.append({"name": os.path.basename(p),
                     "image": base64.b64encode(open(p, "rb").read()).decode()})
    params = {"prompt": args.prompt}
    if args.mode:
        params["mode"] = args.mode
    if args.seed is not None:
        params["seed"] = args.seed
    overrides = {}
    for s in args.sets:
        k, v = s.split("=", 1)
        try:
            overrides[k] = json.loads(v)
        except json.JSONDecodeError:
            overrides[k] = v

    boot = H.wait_for_comfy()
    print(f"comfy ready ({boot:.1f}s)")
    t0 = time.monotonic()
    out = H.handler({"id": f"local-{int(t0)}",
                     "input": {"images": imgs, "params": params, "set": overrides}})
    if "error" in out:
        sys.exit(f"ERROR: {out['error']}")
    for i, img in enumerate(out["images"]):
        dest = pathlib.Path(__file__).parent / f"local-out-{i}-{img['filename']}"
        dest.write_bytes(base64.b64decode(img["data"]))
        print(f"saved {dest}")
    print(f"seed={out['seed']} timings={out['timings']} wall={time.monotonic()-t0:.1f}s")


if __name__ == "__main__":
    main()