comfyui

1#!/usr/bin/env python32"""3run_workflow.py — Inject parameters into a ComfyUI workflow, submit it, monitor4execution, and download outputs.5 6Improvements over v1:7  - Cloud-aware URL routing (handles /api prefix and /history_v2 / /experiment/models renames)8  - API key from CLI flag OR $COMFY_CLOUD_API_KEY env var9  - WebSocket progress monitoring (--ws), with HTTP polling fallback10  - Streaming download (no whole-file buffering — handles GB-size video outputs)11  - Path-traversal-safe output writes12  - Subfolder-aware download paths (no silent overwrites)13  - Retry with exponential backoff on transient errors14  - Status-error correctly classified before "completed: true"15  - Image upload helper (--input-image NAME=PATH)16  - Auto-randomize seed when value is -1 or omitted on a randomize-seed flag17  - Auto-extends timeout heuristically for video workflows18  - Editor-format detection with helpful error19  - Doesn't pollute extra_data.api_key_comfy_org with the cloud auth key20    unless --partner-key is provided (correct semantic per cloud docs)21 22Usage:23    # Local server24    python3 run_workflow.py --workflow workflow_api.json \25        --args '{"prompt": "a cat", "seed": 42}' \26        --output-dir ./outputs27 28    # Cloud server (API key from env var)29    export COMFY_CLOUD_API_KEY="comfyui-xxxxxxx"30    python3 run_workflow.py --workflow workflow_api.json \31        --args '{"prompt": "a cat"}' \32        --host https://cloud.comfy.org \33        --output-dir ./outputs34 35    # With image input (auto-uploads, then references)36    python3 run_workflow.py --workflow img2img.json \37        --input-image image=./photo.png \38        --args '{"prompt": "make it cyberpunk"}'39 40    # WebSocket real-time progress41    python3 run_workflow.py --workflow flux_dev.json \42        --args '{"prompt": "..."}' \43        --ws44 45Stdlib-only by default (Python 3.10+). Will use `requests`/`websocket-client`46if installed for nicer behavior.47"""48 49from __future__ import annotations50 51import argparse52import copy53import json54import sys55import time56from pathlib import Path57from typing import Any58from urllib.parse import urlencode, urlparse59 60# Local import — _common.py sits next to this script.61sys.path.insert(0, str(Path(__file__).resolve().parent))62from _common import (  # noqa: E40263    DEFAULT_LOCAL_HOST, ENV_API_KEY,64    coerce_seed, emit_json, http_get, http_post, http_request,65    is_cloud_host, is_link, log, looks_like_video_workflow,66    media_type_from_filename, new_client_id, resolve_api_key, resolve_url,67    safe_path_join, unwrap_workflow,68)69 70 71# =============================================================================72# Runner73# =============================================================================74 75class WorkflowRunError(Exception):76    """Raised when a workflow run fails (validation, execution, timeout)."""77 78    def __init__(self, status: str, message: str, **details: Any):79        super().__init__(message)80        self.status = status81        self.message = message82        self.details = details83 84    def to_dict(self) -> dict:85        d = {"status": self.status, "error": self.message}86        d.update(self.details)87        return d88 89 90class ComfyRunner:91    def __init__(92        self,93        host: str = DEFAULT_LOCAL_HOST,94        api_key: str | None = None,95        client_id: str | None = None,96        partner_key: str | None = None,97    ):98        self.host = host.rstrip("/")99        self.api_key = api_key100        self.partner_key = partner_key101        self.is_cloud = is_cloud_host(self.host)102        self.client_id = client_id or new_client_id()103 104    @property105    def headers(self) -> dict[str, str]:106        h: dict[str, str] = {}107        if self.api_key:108            h["X-API-Key"] = self.api_key109        return h110 111    def _url(self, path: str) -> str:112        return resolve_url(self.host, path, is_cloud=self.is_cloud)113 114    # ---------- server health ----------115    def check_server(self) -> tuple[bool, dict | None]:116        try:117            r = http_get(self._url("/system_stats"), headers=self.headers, retries=2)118            if r.status == 200:119                try:120                    return True, r.json()121                except Exception:122                    return True, None123            return False, {"http_status": r.status, "body": r.text()[:500]}124        except Exception as e:125            return False, {"error": str(e)}126 127    # ---------- upload ----------128    def upload_image(self, path: Path, *, image_type: str = "input", overwrite: bool = True,129                     endpoint: str = "/upload/image", extra_form: dict | None = None) -> dict:130        """Upload an image file via multipart. Returns server-side ref dict."""131        if not path.exists():132            raise FileNotFoundError(f"input image not found: {path}")133        # Stream the file via a handle to avoid OOM on huge inputs (16MP+ photos).134        with path.open("rb") as fh:135            files = {"image": (path.name, fh)}136            form = {"type": image_type}137            if overwrite:138                form["overwrite"] = "true"139            if extra_form:140                form.update({k: str(v) for k, v in extra_form.items()})141            r = http_request(142                "POST", self._url(endpoint),143                headers=self.headers, files=files, form=form,144                timeout=300, retries=2,145            )146        if r.status != 200:147            raise WorkflowRunError(148                "upload_failed",149                f"Upload of {path.name} failed: HTTP {r.status}",150                body=r.text()[:500],151            )152        try:153            return r.json()154        except Exception:155            return {"name": path.name}156 157    def upload_mask(self, path: Path, original_ref: dict) -> dict:158        """Upload an inpaint mask, linked to a previously uploaded source image.159 160        `original_ref` should be the dict returned by `upload_image()` for the161        source image (or `{"filename": ..., "subfolder": ..., "type": "input"}`).162        """163        return self.upload_image(164            path,165            endpoint="/upload/mask",166            extra_form={167                "subfolder": "clipspace",168                "original_ref": json.dumps(original_ref),169            },170        )171 172    # ---------- submit ----------173    def submit(self, workflow: dict) -> dict:174        payload: dict[str, Any] = {"prompt": workflow, "client_id": self.client_id}175        if self.partner_key:176            payload["extra_data"] = {"api_key_comfy_org": self.partner_key}177 178        r = http_post(self._url("/prompt"), headers=self.headers, json_body=payload, timeout=120)179        try:180            body = r.json()181        except Exception:182            body = {"raw": r.text()[:500]}183        if r.status != 200:184            return {"_http_error": r.status, "body": body}185        return body186 187    # ---------- HTTP polling ----------188    def poll_status(self, prompt_id: str, *, timeout: float = 300.0,189                    initial_interval: float = 1.5, max_interval: float = 8.0) -> dict:190        start = time.time()191        interval = initial_interval192 193        while time.time() - start < timeout:194            if self.is_cloud:195                r = http_get(196                    self._url(f"/job/{prompt_id}/status"),197                    headers=self.headers, retries=2, timeout=30,198                )199                if r.status == 200:200                    try:201                        data = r.json()202                    except Exception:203                        data = {}204                    s = data.get("status")205                    if s == "completed":206                        return {"status": "success", "data": data}207                    if s in {"failed",}:208                        return {"status": "error", "data": data}209                    if s == "cancelled":210                        return {"status": "cancelled", "data": data}211                    # pending / in_progress → continue212                elif r.status == 404:213                    # Cloud sometimes 404s briefly between submit and dispatcher pickup214                    pass215                else:216                    # transient error — retry loop covers it217                    pass218            else:219                # Local: /history/{id} grows once execution completes220                r = http_get(221                    self._url(f"/history/{prompt_id}"),222                    headers=self.headers, retries=2, timeout=30,223                )224                if r.status == 200:225                    try:226                        data = r.json() or {}227                    except Exception:228                        data = {}229                    entry = data.get(prompt_id)230                    if isinstance(entry, dict):231                        st = entry.get("status") or {}232                        # IMPORTANT: check error first — `completed: true` can coexist with errors233                        status_str = st.get("status_str")234                        if status_str == "error":235                            return {"status": "error", "data": entry}236                        if st.get("completed", False):237                            return {"status": "success", "outputs": entry.get("outputs", {})}238                # not in history yet → continue polling239 240            time.sleep(interval)241            interval = min(max_interval, interval * 1.4)242 243        return {"status": "timeout", "elapsed": time.time() - start}244 245    # ---------- WebSocket monitoring ----------246    def monitor_ws(self, prompt_id: str, *, timeout: float = 300.0,247                   on_progress: Any = None) -> dict:248        """Connect to /ws and listen until execution_success / execution_error.249 250        Falls back to HTTP polling if `websocket-client` is not installed.251        Returns same shape as poll_status.252        """253        try:254            import websocket  # type: ignore[import-not-found]255        except ImportError:256            log("websocket-client not installed; falling back to HTTP polling")257            return self.poll_status(prompt_id, timeout=timeout)258 259        # Build WS URL. Preserve any base-path components the user gave us260        # (e.g. http://example.com/comfyui → ws://example.com/comfyui/ws).261        parsed = urlparse(self.host)262        scheme = "wss" if parsed.scheme == "https" else "ws"263        netloc = parsed.netloc264        base_path = parsed.path.rstrip("/")265        ws_url = f"{scheme}://{netloc}{base_path}/ws?clientId={self.client_id}"266        if self.is_cloud and self.api_key:267            ws_url += f"&token={self.api_key}"268 269        outputs: dict[str, Any] = {}270        error_payload: dict[str, Any] | None = None271        success = False272        seen_executed = False273 274        ws = websocket.create_connection(ws_url, timeout=timeout)275        try:276            ws.settimeout(timeout)277            deadline = time.time() + timeout278            while time.time() < deadline:279                msg = ws.recv()280                if isinstance(msg, bytes):281                    # Binary preview frame — ignore for now; ws_monitor.py prints them282                    continue283                try:284                    payload = json.loads(msg)285                except Exception:286                    continue287                mtype = payload.get("type", "")288                mdata = payload.get("data", {}) or {}289 290                # Filter to our job (cloud broadcasts; local filters via client_id)291                pid = mdata.get("prompt_id")292                if pid is not None and pid != prompt_id:293                    continue294 295                if mtype == "progress":296                    if callable(on_progress):297                        on_progress({298                            "type": "progress",299                            "value": mdata.get("value"),300                            "max": mdata.get("max"),301                            "node": mdata.get("node"),302                        })303                elif mtype == "progress_state":304                    if callable(on_progress):305                        on_progress({"type": "progress_state", "nodes": mdata.get("nodes", {})})306                elif mtype == "executing":307                    node = mdata.get("node")308                    if callable(on_progress):309                        on_progress({"type": "executing", "node": node})310                    # When `node` is None on a local server, that signals end-of-run311                    if node is None and not self.is_cloud and seen_executed:312                        success = True313                        break314                elif mtype == "executed":315                    seen_executed = True316                    nid = mdata.get("node")317                    out = mdata.get("output") or {}318                    if nid:319                        outputs[nid] = out320                elif mtype == "notification":321                    if callable(on_progress):322                        on_progress({"type": "notification", "message": mdata.get("value", "")})323                elif mtype == "execution_success":324                    success = True325                    break326                elif mtype == "execution_error":327                    error_payload = mdata328                    break329                elif mtype == "execution_interrupted":330                    error_payload = {"interrupted": True, **mdata}331                    break332        finally:333            try:334                ws.close()335            except Exception:336                pass337 338        if error_payload is not None:339            return {"status": "error", "data": error_payload}340        if success:341            return {"status": "success", "outputs": outputs}342        return {"status": "timeout", "elapsed": timeout}343 344    # ---------- outputs ----------345    def get_outputs(self, prompt_id: str) -> dict:346        if self.is_cloud:347            # Try /jobs/{id} first (returns full job with outputs); fall back to /history_v2348            r = http_get(self._url(f"/jobs/{prompt_id}"), headers=self.headers, retries=2)349            if r.status == 200:350                try:351                    return (r.json() or {}).get("outputs", {}) or {}352                except Exception:353                    pass354            # Fallback355            r = http_get(self._url(f"/history/{prompt_id}"), headers=self.headers, retries=2)356            if r.status == 200:357                try:358                    body = r.json() or {}359                except Exception:360                    body = {}361                if isinstance(body, dict) and prompt_id in body:362                    return body[prompt_id].get("outputs", {}) or {}363                if isinstance(body, dict) and "outputs" in body:364                    return body["outputs"] or {}365            return {}366        # Local367        r = http_get(self._url(f"/history/{prompt_id}"), headers=self.headers, retries=2)368        if r.status != 200:369            return {}370        try:371            body = r.json() or {}372        except Exception:373            return {}374        entry = body.get(prompt_id) or {}375        return entry.get("outputs", {}) or {}376 377    def download_output(378        self, *, filename: str, subfolder: str, file_type: str,379        output_dir: Path, preserve_subfolder: bool = True, overwrite: bool = False,380    ) -> Path:381        """Stream a single output to disk. Path-traversal-safe."""382        params = {"filename": filename, "subfolder": subfolder, "type": file_type}383        url = self._url("/view") + "?" + urlencode(params)384 385        # Compute target path safely. If preserve_subfolder, include subfolder in the386        # local path; otherwise put the file in output_dir flat.387        target_parts: list[str] = []388        if preserve_subfolder and subfolder:389            target_parts.extend(p for p in subfolder.split("/") if p and p not in {".", ".."})390        target_parts.append(filename)391        out_path = safe_path_join(output_dir, *target_parts)392 393        if out_path.exists() and not overwrite:394            stem, suffix = out_path.stem, out_path.suffix395            i = 1396            while True:397                candidate = out_path.with_name(f"{stem}_{i}{suffix}")398                if not candidate.exists():399                    out_path = candidate400                    break401                i += 1402 403        out_path.parent.mkdir(parents=True, exist_ok=True)404 405        # Stream download. Two-step for cloud: get the 302, then fetch signed URL406        # so we don't accidentally send X-API-Key to the storage backend.407        # The HTTP transport already strips X-API-Key on cross-host redirect408        # via _strip_api_key_on_redirect, so a single follow_redirects=True call409        # is safe AND simpler.410        r = http_request(411            "GET", url, headers=self.headers,412            timeout=600, retries=3, follow_redirects=True,413            stream=True, sink=out_path,414        )415        if r.status != 200:416            try:417                if out_path.exists():418                    out_path.unlink()419            except Exception:420                pass421            raise WorkflowRunError(422                "download_failed",423                f"Download of {filename} failed: HTTP {r.status}",424                url=url,425            )426        return out_path427 428    # ---------- queue / cancel ----------429    def cancel(self, prompt_id: str | None = None) -> bool:430        if prompt_id:431            r = http_post(432                self._url("/queue"), headers=self.headers,433                json_body={"delete": [prompt_id]}, retries=1,434            )435            return r.status == 200436        # Interrupt currently running437        r = http_post(self._url("/interrupt"), headers=self.headers, retries=1)438        return r.status == 200439 440 441# =============================================================================442# Schema / parameter injection443# =============================================================================444 445def _inline_schema(workflow: dict) -> dict:446    """Generate schema using the sibling extract_schema module."""447    from extract_schema import extract_schema  # noqa: WPS433448    return extract_schema(workflow)449 450 451def load_schema(schema_path: str | None, workflow: dict) -> dict:452    if schema_path:453        with open(schema_path, encoding="utf-8-sig") as f:454            return json.load(f)455    return _inline_schema(workflow)456 457 458def inject_params(459    workflow: dict, schema: dict, args: dict,460    *, randomize_seed_if_unset: bool = False,461) -> tuple[dict, list[str]]:462    """Inject user args into the workflow. Returns (new_workflow, warnings)."""463    wf = copy.deepcopy(workflow)464    params = schema.get("parameters", {}) or {}465    warnings: list[str] = []466 467    # Auto-randomize seed when it's -1 in args, or when randomize_seed_if_unset468    # and user didn't pass a seed.469    if "seed" in params:470        if "seed" in args and args["seed"] in {None, -1, "-1"}:471            args = dict(args)472            args["seed"] = coerce_seed(args["seed"])473            warnings.append(f"seed=-1 expanded to {args['seed']}")474        elif randomize_seed_if_unset and "seed" not in args:475            args = dict(args)476            args["seed"] = coerce_seed(None)477            warnings.append(f"seed auto-randomized to {args['seed']}")478 479    for name, value in args.items():480        if name not in params:481            warnings.append(f"unknown parameter '{name}' (not in schema), skipping")482            continue483        m = params[name]484        nid, field = m["node_id"], m["field"]485        node = wf.get(nid)486        if not isinstance(node, dict) or "inputs" not in node:487            warnings.append(f"node '{nid}' for parameter '{name}' missing in workflow")488            continue489        # Refuse to overwrite a link with a literal — would silently break wiring490        cur = node["inputs"].get(field)491        if is_link(cur):492            warnings.append(493                f"parameter '{name}' targets {nid}.{field} which is currently a link; "494                f"refusing to overwrite (set the schema to point at the source node instead)"495            )496            continue497        node["inputs"][field] = value498 499    return wf, warnings500 501 502# =============================================================================503# Output download helper504# =============================================================================505 506def download_outputs(507    runner: ComfyRunner, outputs: dict, output_dir: Path,508    *, preserve_subfolder: bool = True, overwrite: bool = False,509) -> list[dict]:510    """Walk the outputs dict and download every file. Cloud uses `video` (singular);511    local uses `videos` (plural). We accept both."""512    output_dir.mkdir(parents=True, exist_ok=True)513    downloaded: list[dict] = []514 515    OUTPUT_KEYS = ("images", "gifs", "videos", "video", "audio", "files", "models", "3d")516 517    for node_id, node_output in (outputs or {}).items():518        if not isinstance(node_output, dict):519            continue520        for key in OUTPUT_KEYS:521            entries = node_output.get(key)522            if not entries:523                continue524            if not isinstance(entries, list):525                entries = [entries]526            for fi in entries:527                if not isinstance(fi, dict):528                    continue529                filename = fi.get("filename") or ""530                if not filename:531                    continue532                subfolder = fi.get("subfolder") or ""533                file_type = fi.get("type") or "output"534                try:535                    out_path = runner.download_output(536                        filename=filename, subfolder=subfolder, file_type=file_type,537                        output_dir=output_dir, preserve_subfolder=preserve_subfolder,538                        overwrite=overwrite,539                    )540                    downloaded.append({541                        "file": str(out_path),542                        "node_id": node_id,543                        "type": media_type_from_filename(filename),544                        "filename": filename,545                        "subfolder": subfolder,546                        "source_type": file_type,547                    })548                except Exception as e:549                    log(f"WARN: failed to download {filename}: {e}")550    return downloaded551 552 553# =============================================================================554# CLI555# =============================================================================556 557def parse_input_image_arg(spec: str) -> tuple[str, Path]:558    """Parse `name=path` (or `path` alone, defaulting to name='image')."""559    if "=" in spec:560        name, path = spec.split("=", 1)561        return name.strip(), Path(path).expanduser()562    return "image", Path(spec).expanduser()563 564 565def main(argv: list[str] | None = None) -> int:566    p = argparse.ArgumentParser(567        description="Run a ComfyUI workflow with parameter injection.",568        formatter_class=argparse.RawDescriptionHelpFormatter,569    )570    p.add_argument("--workflow", required=True, help="Path to workflow API JSON file")571    p.add_argument("--args", default="{}",572                   help="JSON parameters to inject (or `@/path/to/args.json`)")573    p.add_argument("--schema", help="Path to schema JSON (auto-generated if omitted)")574    p.add_argument("--host", default=DEFAULT_LOCAL_HOST, help="ComfyUI server URL")575    p.add_argument("--api-key",576                   help=f"API key for cloud (or set ${ENV_API_KEY} env var)")577    p.add_argument("--partner-key",578                   help="Partner-node API key (extra_data.api_key_comfy_org). "579                        "Required for Flux Pro / Ideogram / etc. Defaults to --api-key if not set.")580    p.add_argument("--output-dir", default="./outputs", help="Directory to save outputs")581    p.add_argument("--timeout", type=int, default=0,582                   help="Max seconds to wait (0=auto: 300 / 900 for video workflows)")583    p.add_argument("--input-image", action="append", default=[],584                   help="Upload local image before running. Format: `name=path` or `path`. "585                        "The `name` becomes the value injected into the matching schema parameter.")586    p.add_argument("--randomize-seed", action="store_true",587                   help="If schema has a 'seed' parameter and --args didn't set one, randomize it")588    p.add_argument("--ws", action="store_true",589                   help="Use WebSocket for real-time progress (requires `websocket-client`)")590    p.add_argument("--no-download", action="store_true", help="Skip downloading outputs")591    p.add_argument("--flat-output", action="store_true",592                   help="Don't preserve server-side subfolder structure when saving outputs")593    p.add_argument("--overwrite", action="store_true",594                   help="Overwrite existing files instead of appending _1, _2, ...")595    p.add_argument("--submit-only", action="store_true",596                   help="Submit and return prompt_id without waiting")597    p.add_argument("--client-id", help="Override generated client_id (UUID)")598    p.add_argument("--use-partner-key-as-auth", action="store_true",599                   help="(Compat) Use --partner-key value as cloud X-API-Key. Don't use unless you know why.")600 601    args = p.parse_args(argv)602 603    # ---- Load workflow ----604    wf_path = Path(args.workflow).expanduser()605    if not wf_path.exists():606        emit_json({"error": f"Workflow file not found: {args.workflow}"})607        return 1608    try:609        with wf_path.open(encoding="utf-8-sig") as f:610            workflow_raw = json.load(f)611        workflow = unwrap_workflow(workflow_raw)612    except ValueError as e:613        emit_json({"error": str(e)})614        return 1615    except json.JSONDecodeError as e:616        emit_json({"error": f"Invalid JSON in workflow file: {e}"})617        return 1618 619    # ---- Parse user args ----620    args_str = args.args621    if args_str.startswith("@"):622        try:623            args_str = Path(args_str[1:]).read_text(encoding="utf-8")624        except OSError as e:625            emit_json({"error": f"Cannot read args file: {e}"})626            return 1627    try:628        user_args = json.loads(args_str) if args_str.strip() else {}629    except json.JSONDecodeError as e:630        emit_json({"error": f"Invalid --args JSON: {e}"})631        return 1632    if not isinstance(user_args, dict):633        emit_json({"error": "--args must be a JSON object"})634        return 1635 636    # ---- Resolve API key ----637    api_key = resolve_api_key(args.api_key)638    partner_key = args.partner_key or None639    if args.use_partner_key_as_auth and not api_key and partner_key:640        api_key = partner_key641 642    # ---- Connect ----643    runner = ComfyRunner(644        host=args.host, api_key=api_key, partner_key=partner_key,645        client_id=args.client_id,646    )647 648    # Server reachability649    ok, info = runner.check_server()650    if not ok:651        emit_json({652            "error": f"Cannot reach server at {args.host}",653            "details": info,654            "hint": (655                "Check `comfy launch --background` is running for local, "656                f"or set ${ENV_API_KEY} for cloud."657            ),658        })659        return 1660 661    # ---- Upload input images ----662    upload_warnings: list[str] = []663    for spec in args.input_image:664        try:665            param_name, path = parse_input_image_arg(spec)666        except Exception as e:667            emit_json({"error": f"Bad --input-image spec '{spec}': {e}"})668            return 1669        try:670            ref = runner.upload_image(path)671        except Exception as e:672            emit_json({"error": f"Upload failed for {path}: {e}"})673            return 1674        # Register as a user arg so inject_params consumes it through the schema675        uploaded_name = ref.get("name") or path.name676        if param_name not in user_args:677            user_args[param_name] = uploaded_name678 679    # ---- Inject params ----680    schema = load_schema(args.schema, workflow)681    workflow, inj_warnings = inject_params(682        workflow, schema, user_args, randomize_seed_if_unset=args.randomize_seed,683    )684    warnings = upload_warnings + inj_warnings685    for w in warnings:686        log(f"WARN: {w}")687 688    # ---- Submit ----689    submit_resp = runner.submit(workflow)690    if "_http_error" in submit_resp:691        emit_json({692            "error": "Submission HTTP error",693            "http_status": submit_resp["_http_error"],694            "body": submit_resp.get("body"),695        })696        return 1697 698    if isinstance(submit_resp.get("error"), dict):699        emit_json({700            "error": "Workflow validation failed",701            "details": submit_resp["error"],702            "node_errors": submit_resp.get("node_errors"),703        })704        return 1705 706    prompt_id = submit_resp.get("prompt_id")707    if not prompt_id:708        emit_json({"error": "No prompt_id in submit response", "response": submit_resp})709        return 1710 711    node_errors = submit_resp.get("node_errors") or {}712    if node_errors:713        emit_json({"error": "Workflow validation failed", "node_errors": node_errors})714        return 1715 716    if args.submit_only:717        emit_json({"status": "submitted", "prompt_id": prompt_id, "warnings": warnings})718        return 0719 720    # ---- Wait ----721    timeout = args.timeout722    if timeout <= 0:723        timeout = 900 if looks_like_video_workflow(workflow) else 300724 725    log(f"Submitted: prompt_id={prompt_id}, waiting (timeout={timeout}s)…")726 727    def _on_progress(evt: dict) -> None:728        t = evt.get("type")729        if t == "progress":730            log(f"  step {evt.get('value')}/{evt.get('max')} on node {evt.get('node')}")731        elif t == "executing":732            node = evt.get("node")733            if node:734                log(f"  executing node {node}")735 736    try:737        if args.ws:738            wait_result = runner.monitor_ws(prompt_id, timeout=timeout, on_progress=_on_progress)739        else:740            wait_result = runner.poll_status(prompt_id, timeout=timeout)741    except KeyboardInterrupt:742        log(f"Interrupted — cancelling job {prompt_id} on server…")743        try:744            runner.cancel(prompt_id)745        except Exception as e:746            log(f"  (cancel request failed: {e})")747        emit_json({748            "status": "interrupted",749            "prompt_id": prompt_id,750            "note": "Ctrl+C received; sent cancellation to server.",751        })752        return 130753 754    if wait_result["status"] == "timeout":755        emit_json({756            "status": "timeout",757            "prompt_id": prompt_id,758            "elapsed": wait_result.get("elapsed"),759            "hint": "Re-run with larger --timeout, or use --submit-only and check later.",760        })761        return 1762    if wait_result["status"] == "error":763        emit_json({"status": "error", "prompt_id": prompt_id, "details": wait_result.get("data")})764        return 1765    if wait_result["status"] == "cancelled":766        emit_json({"status": "cancelled", "prompt_id": prompt_id})767        return 1768 769    # ---- Outputs ----770    outputs = wait_result.get("outputs")771    if not outputs:772        outputs = runner.get_outputs(prompt_id)773 774    if args.no_download:775        emit_json({776            "status": "success", "prompt_id": prompt_id,777            "outputs": outputs, "warnings": warnings,778        })779        return 0780 781    downloaded = download_outputs(782        runner, outputs, Path(args.output_dir).expanduser(),783        preserve_subfolder=not args.flat_output, overwrite=args.overwrite,784    )785 786    emit_json({787        "status": "success",788        "prompt_id": prompt_id,789        "outputs": downloaded,790        "warnings": warnings,791    })792    return 0793 794 795if __name__ == "__main__":796    sys.exit(main())797