File size: 6,112 Bytes
41a10e6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
#!/usr/bin/env python3
"""Call the faceswap-lissie RunPod serverless endpoint SYNCHRONOUSLY (/runsync).

Standalone: python 3.8+, stdlib only. Flux-2-Klein head/face swap
(workflow new-faceswap-lissie-v4): puts the head from --face-image onto the
body/scene of --body-image.

    python faceswap_lissie_sync.py \
        --api-key $RUNPOD_API_KEY \
        --body-image scene.jpg --face-image face.jpg \
        --out ./results

All arguments:
    --api-key           RunPod API key (or set RUNPOD_API_KEY)
    --endpoint-id       RunPod endpoint id     (default: ENDPOINT_DEFAULT below)
    --body-image        image supplying body/scene/pose (required)
    --face-image        image supplying the head/face (required)
    --instruction       head-swap instruction text (sane default baked in)
    --expression-prompt what the QwenVL captioner is asked about the face
    --negative-prompt   negative conditioning text
    --steps             sampler steps (default 5)
    --cfg               guidance (default 1.0)
    --seed              integer seed (default: random, printed back)
    --rmbg-sensitivity / --rmbg-process-res    background-removal tuning
    --workflow          baked workflow name (default new-faceswap-lissie-v4)
    --set NODE.INPUT=VALUE   raw graph override, repeatable
    --workflow-json FILE     full API-format graph passthrough
    --out               output directory (default .)
    --timeout           max seconds (default 600)
"""
import argparse
import base64
import json
import os
import pathlib
import sys
import time
import urllib.error
import urllib.request

DEFAULT_ENDPOINT = "qkeibkeoefcrk3"


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("--body-image", required=True)
    ap.add_argument("--face-image", required=True)
    ap.add_argument("--instruction")
    ap.add_argument("--expression-prompt")
    ap.add_argument("--negative-prompt")
    ap.add_argument("--steps", type=int)
    ap.add_argument("--cfg", type=float)
    ap.add_argument("--seed", type=int)
    ap.add_argument("--rmbg-sensitivity", type=float)
    ap.add_argument("--rmbg-process-res", type=int)
    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")

    params = {"body_image": os.path.basename(args.body_image),
              "face_image": os.path.basename(args.face_image)}
    for cli, p in [("instruction", "instruction"), ("expression_prompt", "expression_prompt"),
                   ("negative_prompt", "negative_prompt"), ("steps", "steps"),
                   ("cfg", "cfg"), ("seed", "seed"),
                   ("rmbg_sensitivity", "rmbg_sensitivity"),
                   ("rmbg_process_res", "rmbg_process_res")]:
        v = getattr(args, cli)
        if v is not None:
            params[p] = v

    payload = {
        "images": [
            {"name": os.path.basename(args.body_image),
             "image": base64.b64encode(open(args.body_image, "rb").read()).decode()},
            {"name": os.path.basename(args.face_image),
             "image": base64.b64encode(open(args.face_image, "rb").read()).decode()},
        ],
        "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

    req = urllib.request.Request(
        f"https://api.runpod.ai/v2/{args.endpoint_id}/runsync",
        data=json.dumps({"input": payload}).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']}")
    out_dir = pathlib.Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    for i, img in enumerate(output["images"]):
        p = out_dir / f"{int(time.time())}-{i}-{img['filename']}"
        p.write_bytes(base64.b64decode(img["data"]))
        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()