| """RunPod serverless handler for ComfyUI workflows. |
| |
| Generic across endpoint images: everything workflow-specific lives in |
| endpoints/<family>/ (API-format workflow JSONs + param maps + endpoint.json). |
| The family baked into an image is selected by $SERVERLESS_ENDPOINT |
| (default: the only directory in endpoints/). |
| |
| Request schema (everything except images/prompt optional): |
| |
| {"input": { |
| "workflow": "qwen-edit-turbo-v4", # default: endpoint.json default_workflow |
| "images": [ # written into ComfyUI's input dir |
| {"name": "main.jpg", "image": "<base64 or data-URI>"}, |
| {"name": "ref.png", "image": "<base64>"} |
| ], |
| "params": { # friendly names from the param map |
| "prompt": "...", |
| "image": "main.jpg", # default: images[0] |
| "image_ref": "ref.png", # default: images[1] if present |
| "mode": "turbo-8", # turbo-4 | turbo-8 | quality |
| "seed": 123, # default: random |
| "input_max_dim": 2048, "output_max_dim": 2560, ... |
| }, |
| "set": {"43.cfg": 1.0}, # raw NODE.INPUT overrides, applied last |
| "workflow_json": { ... } # full API-graph passthrough (advanced; |
| # params other than images/set are ignored) |
| }} |
| |
| Response: |
| |
| {"images": [{"filename": "...", "type": "base64", "data": "..."}], |
| "seed": 123, "prompt_id": "...", |
| "timings": {"queue_s": ..., "execute_s": ..., "total_s": ...}} |
| """ |
|
|
| import base64 |
| import binascii |
| import copy |
| import glob |
| import json |
| import os |
| import random |
| import shutil |
| import sys |
| import time |
| import urllib.error |
| import urllib.request |
|
|
| SERVERLESS_DIR = os.path.dirname(os.path.abspath(__file__)) |
| COMFY_HOST = f"127.0.0.1:{os.environ.get('COMFY_PORT', '7865')}" |
| COMFY_DIR = os.environ.get("COMFY_DIR", "/workspace/ComfyUI") |
| INPUT_DIR = os.path.join(COMFY_DIR, "input") |
| OUTPUT_DIR = os.path.join(COMFY_DIR, "output") |
| BOOT_TIMEOUT = float(os.environ.get("COMFY_BOOT_TIMEOUT", "300")) |
| EXEC_TIMEOUT = float(os.environ.get("COMFY_EXEC_TIMEOUT", "600")) |
| POLL_INTERVAL = 0.25 |
| SEED_INPUTS = ("seed", "noise_seed") |
|
|
|
|
| |
|
|
| def _load_endpoint(): |
| root = os.path.join(SERVERLESS_DIR, "endpoints") |
| name = os.environ.get("SERVERLESS_ENDPOINT") |
| if not name: |
| dirs = sorted(d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))) |
| if len(dirs) != 1: |
| raise RuntimeError(f"set SERVERLESS_ENDPOINT; found {dirs} in {root}") |
| name = dirs[0] |
| base = os.path.join(root, name) |
| with open(os.path.join(base, "endpoint.json")) as fh: |
| cfg = json.load(fh) |
| workflows, params = {}, {} |
| for p in glob.glob(os.path.join(base, "workflows", "*.json")): |
| workflows[os.path.splitext(os.path.basename(p))[0]] = p |
| for p in glob.glob(os.path.join(base, "params", "*.json")): |
| params[os.path.splitext(os.path.basename(p))[0]] = p |
| return name, cfg, workflows, params |
|
|
|
|
| ENDPOINT_NAME, ENDPOINT_CFG, WORKFLOWS, PARAM_MAPS = _load_endpoint() |
|
|
|
|
| |
|
|
| def _http(method, path, payload=None, timeout=30): |
| url = f"http://{COMFY_HOST}{path}" |
| data = json.dumps(payload).encode() if payload is not None else None |
| req = urllib.request.Request(url, data=data, method=method, |
| headers={"Content-Type": "application/json"}) |
| with urllib.request.urlopen(req, timeout=timeout) as r: |
| body = r.read() |
| return json.loads(body) if body else {} |
|
|
|
|
| def wait_for_comfy(timeout=BOOT_TIMEOUT): |
| start = time.monotonic() |
| last_err = None |
| while time.monotonic() - start < timeout: |
| try: |
| _http("GET", "/system_stats", timeout=5) |
| return time.monotonic() - start |
| except Exception as e: |
| last_err = e |
| time.sleep(0.5) |
| log = "" |
| for lf in ("/comfyui.log", os.path.join(COMFY_DIR, "comfyui.log")): |
| if os.path.exists(lf): |
| with open(lf) as fh: |
| log = "".join(fh.readlines()[-40:]) |
| break |
| raise RuntimeError(f"ComfyUI not up after {timeout:.0f}s ({last_err}); log tail:\n{log}") |
|
|
|
|
| |
|
|
| def _decode_image(entry, dest_dir): |
| name = os.path.basename(entry.get("name", "")) |
| if not name: |
| raise ValueError("each images[] entry needs a 'name'") |
| data = entry.get("image", "") |
| if data.startswith("data:"): |
| data = data.split(",", 1)[-1] |
| try: |
| raw = base64.b64decode(data, validate=True) |
| except (binascii.Error, ValueError) as e: |
| raise ValueError(f"images[{name!r}] is not valid base64: {e}") from None |
| os.makedirs(dest_dir, exist_ok=True) |
| with open(os.path.join(dest_dir, name), "wb") as fh: |
| fh.write(raw) |
| return name |
|
|
|
|
| def _apply_param(graph, spec, value): |
| node = graph.get(str(spec["node"])) |
| if node is None: |
| raise ValueError(f"param targets missing node {spec['node']}") |
| inputs = node.setdefault("inputs", {}) |
| if "choices" in spec: |
| if str(value) not in spec["choices"]: |
| raise ValueError(f"value {value!r} not in {sorted(spec['choices'])}") |
| value = spec["choices"][str(value)] |
| if "subkey" in spec: |
| tgt = inputs.get(spec["input"]) |
| if not isinstance(tgt, dict): |
| raise ValueError(f"node {spec['node']}.{spec['input']} is not a dict input") |
| tgt[spec["subkey"]] = value |
| else: |
| inputs[spec["input"]] = value |
|
|
|
|
| def build_graph(job_input, job_id, image_names): |
| """Return (graph, seed, workflow_name).""" |
| if isinstance(job_input.get("workflow_json"), dict): |
| graph = copy.deepcopy(job_input["workflow_json"]) |
| wf_name = "(passthrough)" |
| pmap = {"params": {}, "defaults": {}, "image_slots": []} |
| else: |
| wf_name = job_input.get("workflow") or ENDPOINT_CFG.get("default_workflow") |
| if wf_name not in WORKFLOWS: |
| raise ValueError(f"unknown workflow {wf_name!r}; baked: {sorted(WORKFLOWS)}") |
| with open(WORKFLOWS[wf_name]) as fh: |
| graph = json.load(fh) |
| with open(PARAM_MAPS[wf_name]) as fh: |
| pmap = json.load(fh) |
|
|
| |
| |
| |
| for node in graph.values(): |
| img = node.get("inputs", {}).get("image") |
| if node.get("class_type") == "LoadImage" and img in image_names: |
| node["inputs"]["image"] = f"{job_id}/{img}" |
|
|
| params = dict(pmap.get("defaults", {})) |
| params.update(job_input.get("params") or {}) |
|
|
| |
| |
| |
| |
| slots = pmap.get("image_slots", []) |
| for i, slot in enumerate(slots): |
| if slot not in params and i < len(image_names): |
| params[slot] = image_names[i] |
| for slot, effects in (pmap.get("when_image_absent") or {}).items(): |
| if slot not in params: |
| for k, v in effects.items(): |
| params.setdefault(k, v) |
| for slot in slots: |
| if slot not in params and image_names: |
| params[slot] = image_names[0] |
|
|
| spec_map = pmap.get("params", {}) |
| for key, value in params.items(): |
| spec = spec_map.get(key) |
| if spec is None: |
| raise ValueError(f"unknown param {key!r}; available: {sorted(spec_map)}") |
| if spec.get("type") == "image": |
| if value not in image_names: |
| raise ValueError(f"param {key!r}={value!r} does not match an uploaded image name") |
| value = f"{job_id}/{value}" |
| _apply_param(graph, spec, value) |
|
|
| |
| for target, value in (job_input.get("set") or {}).items(): |
| parts = target.split(".") |
| if len(parts) == 2: |
| spec = {"node": parts[0], "input": parts[1]} |
| elif len(parts) == 3: |
| spec = {"node": parts[0], "input": parts[1], "subkey": parts[2]} |
| else: |
| raise ValueError(f"bad set target {target!r} (want NODE.INPUT or NODE.INPUT.KEY)") |
| _apply_param(graph, spec, value) |
|
|
| |
| |
| |
| |
| seed = params.get("seed") |
| if seed is None: |
| |
| |
| seed = random.randrange(1, 2**31) |
| for node in graph.values(): |
| for name in SEED_INPUTS: |
| if isinstance(node.get("inputs", {}).get(name), int): |
| node["inputs"][name] = seed |
| if "seed" in spec_map: |
| _apply_param(graph, spec_map["seed"], seed) |
|
|
| |
| for node in graph.values(): |
| if "filename_prefix" in node.get("inputs", {}): |
| node["inputs"]["filename_prefix"] = f"{job_id}/out" |
|
|
| return graph, seed, wf_name |
|
|
|
|
| def run_graph(graph): |
| """Queue and wait. Returns (history_entry, queue_s, execute_s).""" |
| t0 = time.monotonic() |
| try: |
| resp = _http("POST", "/prompt", {"prompt": graph, "client_id": "runpod"}) |
| except urllib.error.HTTPError as e: |
| detail = e.read().decode(errors="replace") |
| try: |
| detail = json.dumps(json.loads(detail), indent=2) |
| except ValueError: |
| pass |
| raise RuntimeError(f"workflow validation failed (HTTP {e.code}):\n{detail}") from None |
| prompt_id = resp["prompt_id"] |
| queued = time.monotonic() - t0 |
|
|
| while True: |
| if time.monotonic() - t0 > EXEC_TIMEOUT: |
| _http("POST", "/interrupt", {}) |
| raise RuntimeError(f"execution timed out after {EXEC_TIMEOUT:.0f}s") |
| h = _http("GET", f"/history/{prompt_id}") |
| if prompt_id in h: |
| entry = h[prompt_id] |
| status = entry.get("status", {}) |
| if status.get("status_str") == "error": |
| msgs = [m for m in status.get("messages", []) if m and m[0] == "execution_error"] |
| raise RuntimeError("workflow execution failed:\n" |
| + json.dumps(msgs or status.get("messages", []), indent=2)) |
| return entry, prompt_id, queued, time.monotonic() - t0 |
| time.sleep(POLL_INTERVAL) |
|
|
|
|
| def collect_outputs(entry): |
| out = [] |
| for node_output in entry.get("outputs", {}).values(): |
| for value in node_output.values(): |
| if not isinstance(value, list): |
| continue |
| for item in value: |
| if isinstance(item, dict) and "filename" in item and item.get("type") == "output": |
| sub = item.get("subfolder", "") |
| path = os.path.join(OUTPUT_DIR, sub, item["filename"]) |
| with open(path, "rb") as fh: |
| out.append({"filename": item["filename"], "type": "base64", |
| "data": base64.b64encode(fh.read()).decode()}) |
| |
| |
| |
| try: |
| os.remove(path) |
| except OSError: |
| pass |
| return out |
|
|
|
|
| def handler(job): |
| t_start = time.monotonic() |
| job_id = job.get("id") or f"job-{random.randrange(16**8):08x}" |
| job_id = job_id.replace("/", "_") |
| job_input = job.get("input") or {} |
| job_input_dir = os.path.join(INPUT_DIR, job_id) |
| try: |
| image_names = [_decode_image(e, job_input_dir) |
| for e in (job_input.get("images") or [])] |
| graph, seed, wf_name = build_graph(job_input, job_id, image_names) |
| entry, prompt_id, queue_s, execute_s = run_graph(graph) |
| images = collect_outputs(entry) |
| if not images: |
| return {"error": "workflow finished but produced no output images"} |
| return { |
| "images": images, |
| "seed": seed, |
| "workflow": wf_name, |
| "prompt_id": prompt_id, |
| "timings": {"queue_s": round(queue_s, 3), |
| "execute_s": round(execute_s, 3), |
| "total_s": round(time.monotonic() - t_start, 3)}, |
| } |
| except (ValueError, RuntimeError) as e: |
| return {"error": str(e)} |
| finally: |
| shutil.rmtree(job_input_dir, ignore_errors=True) |
| shutil.rmtree(os.path.join(OUTPUT_DIR, job_id), ignore_errors=True) |
|
|
|
|
| if __name__ == "__main__": |
| boot_wait = wait_for_comfy() |
| print(f"[serverless] endpoint={ENDPOINT_NAME} workflows={sorted(WORKFLOWS)} " |
| f"comfy ready after {boot_wait:.1f}s", flush=True) |
| import runpod |
| runpod.serverless.start({"handler": handler}) |
|
|
