serverless/client/qwen_edit_sync.py

#!/usr/bin/env python3
"""Call the qwen-edit-turbo RunPod serverless endpoint SYNCHRONOUSLY (/runsync).

Standalone: python 3.8+, stdlib only. Blocks until the edit is done (typically
40-65 s warm), saves the returned image(s), prints their paths.

    python qwen_edit_sync.py \
        --api-key $RUNPOD_API_KEY \
        --image person.jpg --ref-image outfit.png \
        --prompt "put the outfit from the second image on the person" \
        --mode turbo-8 --out ./results

All arguments:
    --api-key        RunPod API key (or set RUNPOD_API_KEY env var)
    --endpoint-id    RunPod endpoint id       (default: dom5lwr0o5wq6u)
    --image          main input image (required)
    --ref-image      optional reference image (outfit/style source)
    --prompt         edit instruction (required)
    --mode           turbo-4 | turbo-8 | quality        (default turbo-8)
    --seed           integer seed (default: random; used seed is printed)
    --input-max-dim  max size of the main image fed to the model (default 2048)
    --ref-max-dim    max size of the reference image (default 1024)
    --output-max-dim upscale target for the result (default 2560)
    --lora-skin-fix / --lora-qwen4play      enable the optional loras
    --lora-skin-fix-strength / --lora-qwen4play-strength   (default 1.0)
    --workflow       which baked workflow to run (default qwen-edit-turbo-v4)
    --set NODE.INPUT=VALUE   raw graph override, repeatable (advanced)
    --workflow-json FILE     full API-format graph passthrough (advanced)
    --out            output directory (default .)
    --timeout        max seconds to wait (default 600)
"""
import argparse
import base64
import json
import os
import pathlib
import sys
import time
import urllib.error
import urllib.request

DEFAULT_ENDPOINT = "dom5lwr0o5wq6u"


def build_input(args):
    images, params = [], {"prompt": args.prompt, "mode": args.mode}
    images.append({"name": os.path.basename(args.image),
                   "image": base64.b64encode(open(args.image, "rb").read()).decode()})
    if args.ref_image:
        images.append({"name": os.path.basename(args.ref_image),
                       "image": base64.b64encode(open(args.ref_image, "rb").read()).decode()})
    if args.seed is not None:
        params["seed"] = args.seed
    for cli, param in [("input_max_dim", "input_max_dim"), ("ref_max_dim", "ref_max_dim"),
                       ("output_max_dim", "output_max_dim")]:
        v = getattr(args, cli)
        if v is not None:
            params[param] = v
    if args.lora_skin_fix:
        params["lora_skin_fix"] = True
        params["lora_skin_fix_strength"] = args.lora_skin_fix_strength
    if args.lora_qwen4play:
        params["lora_qwen4play"] = True
        params["lora_qwen4play_strength"] = args.lora_qwen4play_strength

    payload = {"images": images, "params": params}
    if args.workflow:
        payload["workflow"] = args.workflow
    if args.workflow_json:
        payload["workflow_json"] = json.load(open(args.workflow_json))
    overrides = {}
    for s in args.set or []:
        k, v = s.split("=", 1)
        try:
            overrides[k] = json.loads(v)
        except json.JSONDecodeError:
            overrides[k] = v
    if overrides:
        payload["set"] = overrides
    return payload


def save_outputs(output, out_dir):
    out_dir = pathlib.Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    saved = []
    for i, img in enumerate(output.get("images", [])):
        p = out_dir / f"{int(time.time())}-{i}-{img['filename']}"
        p.write_bytes(base64.b64decode(img["data"]))
        saved.append(str(p))
    return saved


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--api-key", default=os.environ.get("RUNPOD_API_KEY"))
    ap.add_argument("--endpoint-id", default=DEFAULT_ENDPOINT)
    ap.add_argument("--image", required=True)
    ap.add_argument("--ref-image")
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--mode", default="turbo-8", choices=["turbo-4", "turbo-8", "quality"])
    ap.add_argument("--seed", type=int)
    ap.add_argument("--input-max-dim", type=int)
    ap.add_argument("--ref-max-dim", type=int)
    ap.add_argument("--output-max-dim", type=int)
    ap.add_argument("--lora-skin-fix", action="store_true")
    ap.add_argument("--lora-skin-fix-strength", type=float, default=1.0)
    ap.add_argument("--lora-qwen4play", action="store_true")
    ap.add_argument("--lora-qwen4play-strength", type=float, default=1.0)
    ap.add_argument("--workflow")
    ap.add_argument("--set", action="append", metavar="NODE.INPUT=VALUE")
    ap.add_argument("--workflow-json")
    ap.add_argument("--out", default=".")
    ap.add_argument("--timeout", type=float, default=600)
    args = ap.parse_args()
    if not args.api_key:
        sys.exit("ERROR: pass --api-key or set RUNPOD_API_KEY")

    req = urllib.request.Request(
        f"https://api.runpod.ai/v2/{args.endpoint_id}/runsync",
        data=json.dumps({"input": build_input(args)}).encode(),
        headers={"Content-Type": "application/json",
                 "Authorization": f"Bearer {args.api_key}"})
    t0 = time.monotonic()
    try:
        with urllib.request.urlopen(req, timeout=args.timeout) as r:
            result = json.load(r)
    except urllib.error.HTTPError as e:
        sys.exit(f"ERROR: HTTP {e.code}: {e.read().decode(errors='replace')[:2000]}")


    # /runsync answers early (IN_QUEUE/IN_PROGRESS) when a job outlives the
    # sync window (~90 s), e.g. during a cold start — fall back to polling
    TERMINAL = ("COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT")
    while result.get("status") not in TERMINAL and result.get("id"):
        if time.monotonic() - t0 > args.timeout:
            sys.exit(f"ERROR: timed out after {args.timeout}s "
                     f"(job {result['id']} status {result.get('status')})")
        time.sleep(3)
        poll = urllib.request.Request(
            f"https://api.runpod.ai/v2/{args.endpoint_id}/status/{result['id']}",
            headers={"Authorization": f"Bearer {args.api_key}"})
        with urllib.request.urlopen(poll, timeout=90) as r:
            result = json.load(r)

    if result.get("status") != "COMPLETED":
        sys.exit(f"ERROR: {json.dumps(result, indent=2)[:3000]}")
    output = result["output"]
    if "error" in output:
        sys.exit(f"ERROR from handler: {output['error']}")
    for p in save_outputs(output, args.out):
        print(p)
    print(f"# seed={output.get('seed')} wall={time.monotonic()-t0:.1f}s "
          f"delay={result.get('delayTime')}ms exec={result.get('executionTime')}ms",
          file=sys.stderr)


if __name__ == "__main__":
    main()