File size: 5,955 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
#!/usr/bin/env python3
"""Call the qwen-edit-nsfw RunPod serverless endpoint SYNCHRONOUSLY (/runsync).

Standalone: python 3.8+, stdlib only. Single-image edit on the
Qwen-Rapid-AIO-NSFW v11.4 checkpoint (workflow qwen-edit-nsfw).

    python qwen_edit_nsfw_sync.py \
        --api-key $RUNPOD_API_KEY \
        --image person.jpg \
        --prompt "replace the outfit with ..." \
        --out ./results

All arguments:
    --api-key         RunPod API key (or set RUNPOD_API_KEY)
    --endpoint-id     RunPod endpoint id      (default: ENDPOINT_DEFAULT below)
    --image           input image (required)
    --prompt          edit instruction (required)
    --negative-prompt negative conditioning text
    --steps           sampler steps (default 8)
    --seed            integer seed (default: random, printed back)
    --megapixels      working resolution in MP (default 1.0)
    --cfg / --eta / --sampler / --scheduler / --denoise    sampler tuning
    --lora-next-scene [--lora-next-scene-strength]         enable next-scene lora
    --workflow        baked workflow name (default qwen-edit-nsfw)
    --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 = "tplcz9jtzihoxa"


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("--prompt", required=True)
    ap.add_argument("--negative-prompt")
    ap.add_argument("--steps", type=int)
    ap.add_argument("--seed", type=int)
    ap.add_argument("--megapixels", type=float)
    ap.add_argument("--cfg", type=float)
    ap.add_argument("--eta", type=float)
    ap.add_argument("--sampler")
    ap.add_argument("--scheduler")
    ap.add_argument("--denoise", type=float)
    ap.add_argument("--lora-next-scene", action="store_true")
    ap.add_argument("--lora-next-scene-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")

    params = {"prompt": args.prompt}
    for cli, p in [("negative_prompt", "negative_prompt"), ("steps", "steps"),
                   ("seed", "seed"), ("megapixels", "megapixels"), ("cfg", "cfg"),
                   ("eta", "eta"), ("sampler", "sampler"), ("scheduler", "scheduler"),
                   ("denoise", "denoise")]:
        v = getattr(args, cli)
        if v is not None:
            params[p] = v
    if args.lora_next_scene:
        params["lora_next_scene"] = True
        params["lora_next_scene_strength"] = args.lora_next_scene_strength

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