processor.py

"""
ComfyUI Generate — Modly process extension (stdlib only, no dependencies).

Runs an API-format ComfyUI workflow on a local server:
  1. loads the workflow JSON (rejects UI-format exports with a clear message)
  2. injects the positive/negative prompts into the right CLIPTextEncode nodes
     (traced by following which sampler slot consumes each encoder)
  3. sets every seed (random by default; a fixed `seed` param >= 0 makes runs
     reproducible — identical inputs are then served from ComfyUI's cache)
  4. POST /prompt -> poll /history/{id} -> GET /view -> saves the image

Protocol: reads one JSON line from stdin, writes JSON lines to stdout.
  stdin : { input: { texts?, text? }, params, workspaceDir, tempDir }
  stdout: { type: "progress"|"log"|"done"|"error", ... }
"""
from __future__ import annotations

import json
import random
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from pathlib import Path

POLL_INTERVAL_S = 1.0
MAX_WAIT_S      = 30 * 60

TEXT_ENCODER_PREFIX = "CLIPTextEncode"
WRAPPER_RE          = re.compile(r"Conditioning|Guidance|Guider", re.IGNORECASE)


def emit(obj: dict) -> None:
    print(json.dumps(obj), flush=True)


def progress(pct: int, label: str) -> None:
    emit({"type": "progress", "percent": pct, "label": label})


def error(msg: str) -> None:
    emit({"type": "error", "message": msg})


def http_json(url: str, payload: dict | None = None, timeout: float = 30.0) -> dict:
    data = json.dumps(payload).encode("utf-8") if payload is not None else None
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as res:
        return json.loads(res.read().decode("utf-8"))


def is_node_ref(value) -> bool:
    return isinstance(value, list) and len(value) == 2 and isinstance(value[0], str)


def trace_text_role(prompt: dict, start_id: str) -> str | None:
    """Follows an encoder's output through conditioning wrappers until a consumer
    plugs it into a `positive` or `negative` slot (e.g. KSampler, or a
    FluxGuidance chain), so we know which prompt to inject."""
    frontier, seen = {start_id}, {start_id}
    for _ in range(5):
        nxt = set()
        for consumer_id, consumer in prompt.items():
            for key, value in consumer.get("inputs", {}).items():
                if not is_node_ref(value) or value[0] not in frontier:
                    continue
                if key == "positive":
                    return "positive"
                if key == "negative":
                    return "negative"
                if WRAPPER_RE.search(consumer.get("class_type", "")) and consumer_id not in seen:
                    seen.add(consumer_id)
                    nxt.add(consumer_id)
        if not nxt:
            return None
        frontier = nxt
    return None


def inject_text(node: dict, text: str) -> None:
    inputs = node.setdefault("inputs", {})
    if "text" in inputs:
        inputs["text"] = text
    if "text_g" in inputs:
        inputs["text_g"] = text
    if "text_l" in inputs:
        inputs["text_l"] = text


def main() -> None:
    data = json.loads(sys.stdin.readline())
    input_data    = data.get("input", {}) or {}
    params        = data.get("params", {}) or {}
    workspace_dir = data.get("workspaceDir", "")

    server_url = (params.get("server_url") or "http://127.0.0.1:8188").rstrip("/")
    wf_dir  = (params.get("workflow_dir") or "").strip()
    wf_file = (params.get("workflow_file") or "").strip()
    workflow_path = str(Path(wf_dir) / wf_file) if wf_dir and wf_file else ""

    try:
        fixed_seed = int(params.get("seed", -1))
    except (TypeError, ValueError):
        fixed_seed = -1

    # Each prompt comes strictly from its own handle (input-0 → positive,
    # input-1 → negative) so they can never bleed into each other. The generic
    # `text` field is only a fallback for a single, unpaired text input.
    texts    = input_data.get("texts") or []
    positive = (texts[0] if len(texts) > 0 else None) or None
    negative = (texts[1] if len(texts) > 1 else None) or None
    if positive is None and negative is None:
        positive = input_data.get("text")

    if not workflow_path or not Path(workflow_path).is_file():
        error(f"ComfyUI: workflow file not found: {workflow_path or '(not set)'}")
        return

    # ── Load + validate the workflow ──────────────────────────────────────────
    raw = json.loads(Path(workflow_path).read_text(encoding="utf-8"))
    if isinstance(raw, dict) and isinstance(raw.get("nodes"), list) and "links" in raw:
        error("This is a UI-format export. In ComfyUI use Workflow > Export (API) instead.")
        return
    prompt = {k: v for k, v in raw.items() if isinstance(v, dict) and "class_type" in v}
    if not prompt:
        error("No executable nodes found — export the workflow in API format.")
        return

    # ── Inject prompts + set seeds ────────────────────────────────────────────
    for node_id, node in prompt.items():
        cls = node.get("class_type", "")
        if cls.startswith(TEXT_ENCODER_PREFIX):
            role = trace_text_role(prompt, node_id)
            if role == "positive" and positive:
                inject_text(node, positive)
            elif role == "negative" and negative:
                inject_text(node, negative)
        for key in ("seed", "noise_seed"):
            if isinstance(node.get("inputs", {}).get(key), (int, float)):
                node["inputs"][key] = fixed_seed if fixed_seed >= 0 else random.randint(0, 2**48)

    # ── Submit ────────────────────────────────────────────────────────────────
    progress(5, "Submitting to ComfyUI…")
    try:
        res = http_json(f"{server_url}/prompt", {"prompt": prompt, "client_id": str(uuid.uuid4())})
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")[:400]
        error(f"ComfyUI rejected the workflow (is it exported in API format?): {detail}")
        return
    except urllib.error.URLError:
        error(f"ComfyUI server not reachable at {server_url} — is it running?")
        return
    prompt_id = res["prompt_id"]

    # ── Poll history ──────────────────────────────────────────────────────────
    progress(50, "Generating…")
    started = time.time()
    entry = None
    while True:
        if time.time() - started > MAX_WAIT_S:
            error("ComfyUI: timed out after 30 minutes")
            return
        time.sleep(POLL_INTERVAL_S)
        history = http_json(f"{server_url}/history/{prompt_id}")
        e = history.get(prompt_id, history)
        if not isinstance(e, dict):
            continue
        status = e.get("status") or {}
        if status.get("status_str") == "error":
            error("ComfyUI: workflow execution failed (check the ComfyUI console)")
            return
        if e.get("outputs") or status.get("completed"):
            entry = e
            break

    # ── Download the first image output ───────────────────────────────────────
    progress(90, "Downloading output…")
    ref = None
    for node_output in (entry.get("outputs") or {}).values():
        for value in node_output.values():
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, dict) and "filename" in item:
                        ref = item
                        break
            if ref:
                break
        if ref:
            break
    if not ref:
        if fixed_seed >= 0:
            error(
                f"ComfyUI produced no outputs — this run (seed {fixed_seed}) was likely "
                "served from ComfyUI's cache. Change the seed or a prompt, or set Seed to -1."
            )
        else:
            error("ComfyUI produced no outputs — check that the workflow has a Save Image node")
        return

    query = urllib.parse.urlencode({
        "filename":  ref["filename"],
        "subfolder": ref.get("subfolder", ""),
        "type":      ref.get("type", "output"),
    })
    out_dir = Path(workspace_dir) / "Workflows"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / f"comfyui-{int(time.time() * 1000)}-{ref['filename'].replace('/', '_')}"
    with urllib.request.urlopen(f"{server_url}/view?{query}", timeout=120) as res:
        out_path.write_bytes(res.read())

    progress(100, "Done")
    emit({"type": "done", "result": {"filePath": str(out_path)}})


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:  # noqa: BLE001 — surface anything to the UI
        error(f"ComfyUI extension failed: {exc}")