foodlinkk-command-center/tools-api/app/comfyui.py

"""ComfyUI client — HD generation with resilient progress tracking."""
from __future__ import annotations

import asyncio
import json
import logging
import os
import random
import uuid
from typing import Any, Optional
from urllib.parse import urlencode

import httpx

log = logging.getLogger("tools-api.comfyui")

COMFYUI_URL = os.getenv("COMFYUI_URL", "http://10.4.7.18:8188").rstrip("/")
CHECKPOINT = os.getenv("COMFYUI_CHECKPOINT", "v1-5-pruned-emaonly.safetensors")
DEFAULT_STEPS = int(os.getenv("COMFYUI_STEPS", "15"))
POLL_INTERVAL = float(os.getenv("COMFYUI_POLL_INTERVAL", "2.0"))
MAX_WAIT = float(os.getenv("COMFYUI_MAX_WAIT", "1800"))
WS_IDLE_TIMEOUT = float(os.getenv("COMFYUI_WS_IDLE", "90"))

QUALITY_PRESETS: dict[str, dict[str, Any]] = {
    "fast": {"width": 512, "height": 512, "steps": 15, "label": "Snel (512px)"},
    "hd": {"width": 1024, "height": 1024, "steps": 28, "label": "HD (1024px)"},
    "ultra": {"width": 1024, "height": 1024, "steps": 35, "label": "Ultra HD (1024px, 35 steps)"},
}

_jobs: dict[str, dict[str, Any]] = {}


def resolve_quality(
    quality: str | None = None,
    width: int | None = None,
    height: int | None = None,
    steps: int | None = None,
) -> tuple[int, int, int, str]:
    q = (quality or "hd").lower()
    preset = QUALITY_PRESETS.get(q, QUALITY_PRESETS["hd"])
    w = width or preset["width"]
    h = height or preset["height"]
    s = steps or preset["steps"]
    label = preset["label"]
    return w, h, s, label


def get_job(prompt_id: str) -> dict[str, Any] | None:
    return _jobs.get(prompt_id)


def build_workflow(
    prompt: str,
    negative: str = "blurry, low quality, watermark, text, ugly, deformed",
    width: int = 1024,
    height: int = 1024,
    steps: int = DEFAULT_STEPS,
    seed: Optional[int] = None,
) -> dict[str, Any]:
    seed = seed if seed is not None else random.randint(1, 2**31 - 1)
    return {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": seed,
                "steps": steps,
                "cfg": 7.5,
                "sampler_name": "euler",
                "scheduler": "normal",
                "denoise": 1.0,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0],
            },
        },
        "4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CHECKPOINT}},
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": width, "height": height, "batch_size": 1},
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": prompt, "clip": ["4", 1]},
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": negative, "clip": ["4", 1]},
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]},
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "foodlinkk", "images": ["8", 0]},
        },
    }


def view_url(filename: str, subfolder: str = "", img_type: str = "output") -> str:
    params = urlencode({"filename": filename, "type": img_type, "subfolder": subfolder})
    return f"{COMFYUI_URL}/view?{params}"


def _new_job(prompt_id: str, prompt: str, width: int, height: int, steps: int, quality: str) -> None:
    _jobs[prompt_id] = {
        "prompt_id": prompt_id,
        "status": "queued",
        "percent": 0,
        "step": 0,
        "max_step": steps,
        "node": None,
        "message": "In wachtrij bij ComfyUI…",
        "prompt": prompt[:500],
        "width": width,
        "height": height,
        "steps": steps,
        "quality": quality,
        "events": [],
        "result": None,
        "error": None,
    }


def _append_event(prompt_id: str, message: str) -> None:
    job = _jobs.get(prompt_id)
    if not job:
        return
    job["message"] = message
    events: list[str] = job.setdefault("events", [])
    if not events or events[-1] != message:
        events.append(message)
    if len(events) > 50:
        del events[: len(events) - 50]


async def submit_prompt(workflow: dict[str, Any]) -> tuple[str, str]:
    client_id = str(uuid.uuid4())
    async with httpx.AsyncClient(timeout=30.0) as client:
        resp = await client.post(
            f"{COMFYUI_URL}/prompt",
            json={"prompt": workflow, "client_id": client_id},
        )
        resp.raise_for_status()
        data = resp.json()
    if data.get("node_errors"):
        raise RuntimeError(f"ComfyUI node errors: {data['node_errors']}")
    return data["prompt_id"], client_id


async def _prompt_in_history(prompt_id: str) -> bool:
    async with httpx.AsyncClient(timeout=20.0) as client:
        resp = await client.get(f"{COMFYUI_URL}/history/{prompt_id}")
        if resp.status_code == 200 and prompt_id in resp.json():
            return True
    return False


async def _prompt_in_queue(prompt_id: str) -> bool:
    async with httpx.AsyncClient(timeout=20.0) as client:
        resp = await client.get(f"{COMFYUI_URL}/queue")
        if resp.status_code != 200:
            return False
        data = resp.json()
        for bucket in ("queue_running", "queue_pending"):
            for item in data.get(bucket) or []:
                if isinstance(item, (list, tuple)) and len(item) > 1 and item[1] == prompt_id:
                    return True
    return False


async def _is_still_running(prompt_id: str) -> bool:
    if await _prompt_in_history(prompt_id):
        return False
    return await _prompt_in_queue(prompt_id)


async def wait_for_output(prompt_id: str) -> dict[str, Any]:
    deadline = asyncio.get_event_loop().time() + MAX_WAIT
    tick = 0
    async with httpx.AsyncClient(timeout=30.0) as client:
        while asyncio.get_event_loop().time() < deadline:
            resp = await client.get(f"{COMFYUI_URL}/history/{prompt_id}")
            if resp.status_code == 200:
                hist = resp.json()
                if prompt_id in hist:
                    outputs = hist[prompt_id].get("outputs") or {}
                    for node_out in outputs.values():
                        images = node_out.get("images") or []
                        if images:
                            img = images[0]
                            return {
                                "filename": img["filename"],
                                "subfolder": img.get("subfolder", ""),
                                "type": img.get("type", "output"),
                            }
            tick += 1
            if tick % 15 == 0:
                _append_event(prompt_id, "ComfyUI CPU render duurt even — nog bezig…")
            await asyncio.sleep(POLL_INTERVAL)
    raise TimeoutError(f"ComfyUI generation timed out after {int(MAX_WAIT)}s")


async def _track_ws(client_id: str, prompt_id: str) -> None:
    try:
        import websockets
    except ImportError:
        _append_event(prompt_id, "Polling modus (geen websocket)")
        return

    ws_url = COMFYUI_URL.replace("https://", "wss://").replace("http://", "ws://") + f"/ws?clientId={client_id}"
    try:
        async with websockets.connect(ws_url, ping_interval=30, ping_timeout=60, close_timeout=10) as ws:
            finished = False
            while not finished:
                try:
                    raw = await asyncio.wait_for(ws.recv(), timeout=WS_IDLE_TIMEOUT)
                except asyncio.TimeoutError:
                    if await _prompt_in_history(prompt_id):
                        finished = True
                        break
                    if await _is_still_running(prompt_id):
                        _append_event(prompt_id, "Nog bezig op CPU (geen WS update)…")
                        continue
                    break

                data = json.loads(raw)
                msg_type = data.get("type")
                payload = data.get("data") or {}
                pid = payload.get("prompt_id")
                if pid not in (None, prompt_id):
                    continue

                if msg_type == "execution_start":
                    _jobs[prompt_id]["status"] = "running"
                    _append_event(prompt_id, "ComfyUI gestart")

                elif msg_type == "progress":
                    val = int(payload.get("value") or 0)
                    mx = int(payload.get("max") or 1)
                    pct = int(100 * val / mx) if mx else 0
                    _jobs[prompt_id].update(
                        status="running",
                        percent=pct,
                        step=val,
                        max_step=mx,
                        node=payload.get("node"),
                    )
                    _append_event(prompt_id, f"KSampler {val}/{mx} ({pct}%)")

                elif msg_type == "executing":
                    node = payload.get("node")
                    if node is None:
                        _jobs[prompt_id]["status"] = "finishing"
                        _append_event(prompt_id, "Render klaar — opslaan…")
                        finished = True
                    else:
                        _jobs[prompt_id]["node"] = node
                        _append_event(prompt_id, f"Node {node}")

                elif msg_type == "execution_error":
                    err = payload.get("exception_message") or "ComfyUI execution error"
                    raise RuntimeError(str(err))
    except Exception as exc:
        log.warning("WS tracking ended for %s: %s — falling back to poll", prompt_id, exc)
        if await _prompt_in_history(prompt_id):
            return
        if await _is_still_running(prompt_id):
            _append_event(prompt_id, "Voortgang via polling (WS verbroken)")
            return
        raise


async def _run_job(
    prompt_id: str,
    client_id: str,
    prompt: str,
    width: int,
    height: int,
    steps: int,
    quality: str,
    seed: Optional[int],
) -> None:
    try:
        try:
            await _track_ws(client_id, prompt_id)
        except Exception as ws_exc:
            log.warning("WS phase issue %s: %s", prompt_id, ws_exc)
            if not await _is_still_running(prompt_id) and not await _prompt_in_history(prompt_id):
                raise

        img = await wait_for_output(prompt_id)
        result = {
            "prompt_id": prompt_id,
            "filename": img["filename"],
            "subfolder": img.get("subfolder", ""),
            "type": img.get("type", "output"),
            "image_url": view_url(img["filename"], img.get("subfolder", ""), img.get("type", "output")),
            "prompt": prompt,
            "width": width,
            "height": height,
            "steps": steps,
            "quality": quality,
        }
        _jobs[prompt_id].update(status="done", percent=100, result=result, message="Klaar!")
        _append_event(prompt_id, "Afbeelding klaar")
        log.info("ComfyUI done %s (%dx%d)", prompt_id, width, height)
    except Exception as exc:
        log.exception("ComfyUI job failed %s", prompt_id)
        if await _prompt_in_history(prompt_id):
            try:
                img = await wait_for_output(prompt_id)
                result = {
                    "prompt_id": prompt_id,
                    "filename": img["filename"],
                    "subfolder": img.get("subfolder", ""),
                    "type": img.get("type", "output"),
                    "image_url": view_url(img["filename"], img.get("subfolder", ""), img.get("type", "output")),
                    "prompt": prompt,
                    "width": width,
                    "height": height,
                    "steps": steps,
                    "quality": quality,
                }
                _jobs[prompt_id].update(status="done", percent=100, result=result, message="Klaar!")
                return
            except Exception:
                pass
        _jobs[prompt_id].update(status="error", error=str(exc), message=str(exc))
        _append_event(prompt_id, f"Fout: {exc}")


async def start_generation(
    prompt: str,
    *,
    negative: str = "blurry, low quality, watermark, text, ugly, deformed",
    quality: str = "hd",
    width: int | None = None,
    height: int | None = None,
    steps: int | None = None,
    seed: Optional[int] = None,
) -> dict[str, Any]:
    w, h, s, label = resolve_quality(quality, width, height, steps)
    workflow = build_workflow(prompt, negative=negative, width=w, height=h, steps=s, seed=seed)
    prompt_id, client_id = await submit_prompt(workflow)
    _new_job(prompt_id, prompt, w, h, s, quality)
    _append_event(prompt_id, f"Gestart — {label}")
    asyncio.create_task(_run_job(prompt_id, client_id, prompt, w, h, s, quality, seed))
    return {
        "prompt_id": prompt_id,
        "client_id": client_id,
        "quality": quality,
        "width": w,
        "height": h,
        "steps": s,
        "quality_label": label,
    }


async def generate_image(
    prompt: str,
    width: int = 1024,
    height: int = 1024,
    steps: int = DEFAULT_STEPS,
    seed: Optional[int] = None,
    quality: str = "hd",
) -> dict[str, Any]:
    if quality and quality != "custom":
        width, height, steps, _ = resolve_quality(quality, width, height, steps)
    started = await start_generation(
        prompt,
        quality="custom",
        width=width,
        height=height,
        steps=steps,
        seed=seed,
    )
    prompt_id = started["prompt_id"]
    deadline = asyncio.get_event_loop().time() + MAX_WAIT
    while asyncio.get_event_loop().time() < deadline:
        job = _jobs.get(prompt_id) or {}
        if job.get("status") == "done" and job.get("result"):
            return job["result"]
        if job.get("status") == "error":
            raise RuntimeError(job.get("error") or "Generation failed")
        await asyncio.sleep(POLL_INTERVAL)
    raise TimeoutError(f"ComfyUI generation timed out after {int(MAX_WAIT)}s")


async def fetch_image_bytes(filename: str, subfolder: str = "", img_type: str = "output") -> bytes:
    params = {"filename": filename, "type": img_type, "subfolder": subfolder}
    async with httpx.AsyncClient(timeout=120.0) as client:
        resp = await client.get(f"{COMFYUI_URL}/view", params=params)
        resp.raise_for_status()
        return resp.content