File size: 2,194 Bytes
b83a9c1 41a10e6 b83a9c1 41a10e6 b83a9c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | #!/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()
|