mcp-comfyui/server.py

"""ComfyUI FLUX.1-schnell MCP server — async queue edition.

Tools:
  generate_image     — txt2img; submits job and returns job_id immediately
  img2img            — img2img for one or more input images; one job per image
  get_job_status     — poll status + retrieve completed images inline
  list_queue         — list recent jobs in the in-process registry
  cancel_job         — cancel a queued/running job (confirm with user first)
  list_recent_images — show n most recent completed images for reference/chaining

Job chaining: pass "job:<job_id>" in the img2img images list to use a previous
output as input without routing bytes through the LLM.
"""
import base64
import copy
import io
import json
import os
import random
import threading
import time
import uuid
from datetime import datetime, timezone

import boto3
import httpx
from botocore.config import Config
from mcp.server.fastmcp import FastMCP, Image

COMFY = os.environ.get("COMFYUI_URL", "http://comfyui:8188")

# S3 upload (optional — skipped if credentials absent)
_S3_ENDPOINT   = os.environ.get("CHAT_ARTIFACTS_ENDPOINT", "")
_S3_ACCESS     = os.environ.get("CHAT_ARTIFACTS_ACCESS_KEY", "")
_S3_SECRET     = os.environ.get("CHAT_ARTIFACTS_SECRET_KEY", "")
_S3_BUCKET     = os.environ.get("CHAT_ARTIFACTS_BUCKET", "chat-artifacts")
_S3_PUBLIC_BASE = "https://chat-artifacts.s3.nuclide.systems"
_S3_ENABLED    = bool(_S3_ENDPOINT and _S3_ACCESS and _S3_SECRET)


def _s3_upload_png(png: bytes, job_id: str) -> str | None:
    if not _S3_ENABLED:
        return None
    key = datetime.now(timezone.utc).strftime("%Y-%m-%d/") + f"comfyui-{job_id}.png"
    try:
        boto3.client(
            "s3",
            endpoint_url=_S3_ENDPOINT,
            aws_access_key_id=_S3_ACCESS,
            aws_secret_access_key=_S3_SECRET,
            region_name="garage",
            config=Config(
                signature_version="s3v4",
                request_checksum_calculation="when_required",
                response_checksum_validation="when_required",
            ),
        ).put_object(Bucket=_S3_BUCKET, Key=key, Body=png, ContentType="image/png")
        return f"{_S3_PUBLIC_BASE}/{key}"
    except Exception:
        return None
_TXT2IMG_PATH = os.environ.get("WORKFLOW_PATH", "/app/flux-gguf-api.json")
_IMG2IMG_PATH = os.environ.get("IMG2IMG_WORKFLOW_PATH", "/app/flux-img2img-api.json")
POLL_INTERVAL = 3
GEN_DEADLINE = int(os.environ.get("GEN_DEADLINE", 480))   # seconds before blocking tools give up

with open(_TXT2IMG_PATH) as fh:
    _BASE_TXT2IMG = json.load(fh)
with open(_IMG2IMG_PATH) as fh:
    _BASE_IMG2IMG = json.load(fh)

mcp = FastMCP("comfyui-flux", host="0.0.0.0", port=8000, stateless_http=True)

# ---------------------------------------------------------------------------
# Job registry: job_id → {comfy_pid, status, kind, prompt, submitted_at,
#                          result_png, output_filename, error}
# ---------------------------------------------------------------------------
_jobs: dict[str, dict] = {}
_lock = threading.Lock()


def _register_job(comfy_pid: str, prompt: str, kind: str) -> str:
    job_id = f"{kind[:3]}-{uuid.uuid4().hex[:8]}"
    with _lock:
        _jobs[job_id] = {
            "comfy_pid": comfy_pid,
            "status": "running",
            "kind": kind,
            "prompt": prompt,
            "submitted_at": time.time(),
            "result_png": None,
            "output_filename": None,
            "error": None,
        }
    return job_id


def _poll_loop():
    while True:
        time.sleep(POLL_INTERVAL)
        with _lock:
            pending = [(jid, j["comfy_pid"]) for jid, j in _jobs.items() if j["status"] == "running"]
        for jid, cpid in pending:
            try:
                with httpx.Client(timeout=10) as c:
                    hist = c.get(f"{COMFY}/history/{cpid}").json()
                if cpid not in hist:
                    continue
                entry = hist[cpid]
                st = entry.get("status", {})
                if st.get("status_str") == "error":
                    with _lock:
                        _jobs[jid]["status"] = "error"
                        _jobs[jid]["error"] = str(st.get("messages", ""))[:500]
                    continue
                imgs = [im for node in entry.get("outputs", {}).values() for im in node.get("images", [])]
                if imgs:
                    im = imgs[0]
                    with httpx.Client(timeout=120) as c:
                        png = c.get(
                            f"{COMFY}/view",
                            params={"filename": im["filename"], "subfolder": im.get("subfolder", ""), "type": im.get("type", "output")},
                        ).content
                    s3_url = _s3_upload_png(png, jid)
                    with _lock:
                        _jobs[jid]["status"] = "done"
                        _jobs[jid]["result_png"] = png
                        _jobs[jid]["output_filename"] = im["filename"]
                        _jobs[jid]["s3_url"] = s3_url
                elif st.get("completed"):
                    with _lock:
                        _jobs[jid]["status"] = "error"
                        _jobs[jid]["error"] = "Completed with no output images"
            except Exception:
                pass


threading.Thread(target=_poll_loop, daemon=True).start()


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _upload_image(img_bytes: bytes, filename: str = "input.png") -> str:
    """Upload image bytes to ComfyUI /upload/image, return the filename assigned."""
    with httpx.Client(timeout=60) as c:
        r = c.post(
            f"{COMFY}/upload/image",
            files={"image": (filename, io.BytesIO(img_bytes), "image/png")},
            data={"type": "input", "overwrite": "true"},
        )
        r.raise_for_status()
    return r.json()["name"]


def _resolve_image(src: str) -> bytes:
    """Resolve an image source to raw bytes.

    Accepts:
      job:<job_id>              — output of a previous completed job
      https://... / http://...  — downloaded URL
      data:image/...;base64,...  — inline base64
    """
    if src.startswith("job:"):
        jid = src[4:].strip()
        with _lock:
            job = _jobs.get(jid)
        if not job:
            raise ValueError(f"Unknown job: {jid}")
        if job["status"] != "done":
            raise ValueError(f"Job {jid} not done yet (status: {job['status']})")
        return job["result_png"]
    if src.startswith("data:"):
        _, payload = src.split(",", 1)
        return base64.b64decode(payload)
    with httpx.Client(timeout=60, follow_redirects=True) as c:
        r = c.get(src)
        r.raise_for_status()
    return r.content


def _submit_txt2img(prompt: str, width: int, height: int, steps: int, seed: int) -> str:
    wf = copy.deepcopy(_BASE_TXT2IMG)
    wf["4"]["inputs"]["text"] = prompt
    wf["6"]["inputs"]["width"] = width
    wf["6"]["inputs"]["height"] = height
    wf["8"]["inputs"]["steps"] = steps
    wf["8"]["inputs"]["seed"] = seed
    with httpx.Client(timeout=30) as c:
        r = c.post(f"{COMFY}/prompt", json={"prompt": wf})
        r.raise_for_status()
    return r.json()["prompt_id"]


def _submit_img2img(prompt: str, img_bytes: bytes, strength: float, steps: int, seed: int) -> str:
    fname = _upload_image(img_bytes, f"i2i_{uuid.uuid4().hex[:6]}.png")
    wf = copy.deepcopy(_BASE_IMG2IMG)
    wf["4"]["inputs"]["text"] = prompt
    wf["11"]["inputs"]["image"] = fname
    wf["8"]["inputs"]["steps"] = steps
    wf["8"]["inputs"]["seed"] = seed
    wf["8"]["inputs"]["denoise"] = max(0.05, min(1.0, strength))
    with httpx.Client(timeout=30) as c:
        r = c.post(f"{COMFY}/prompt", json={"prompt": wf})
        r.raise_for_status()
    return r.json()["prompt_id"]


# ---------------------------------------------------------------------------
# MCP tools
# ---------------------------------------------------------------------------

def _wait_for_job(job_id: str) -> list:
    """Block until the job is done/error/cancelled, or GEN_DEADLINE elapses.
    Returns an MCP content list (text + optional Image).
    """
    deadline = time.time() + GEN_DEADLINE
    while time.time() < deadline:
        with _lock:
            job = dict(_jobs.get(job_id, {}))
        elapsed = int(time.time() - job.get("submitted_at", time.time()))
        if job.get("status") == "done":
            url_line = f"\nS3: {job['s3_url']}" if job.get("s3_url") else ""
            return [
                f"**{job_id}** — done in {elapsed}s\nPrompt: {job['prompt'][:120]}{url_line}",
                Image(data=job["result_png"], format="png"),
            ]
        if job.get("status") == "error":
            return [f"**{job_id}** — generation failed after {elapsed}s: {job.get('error','unknown')}"]
        if job.get("status") == "cancelled":
            return [f"**{job_id}** — cancelled after {elapsed}s"]
        time.sleep(POLL_INTERVAL)
    # Deadline exceeded — return job_id for manual follow-up
    return [
        f"**{job_id}** — still generating after {GEN_DEADLINE}s.\n"
        f"Call `get_job_status([\"{job_id}\"])` to retrieve the image when ready."
    ]


@mcp.tool()
def generate_image(
    prompt: str,
    width: int = 1024,
    height: int = 1024,
    steps: int = 4,
    seed: int = 0,
) -> list:
    """Generate an image from a text prompt. Waits for the result and returns it inline.

    Generation takes 90-300 s; this tool blocks until the image is ready — no
    follow-up call needed. The image is returned directly in this response.

    FLUX renders legible in-image text well — put quoted text in the prompt.
    4 steps is optimal for schnell. seed=0 picks a random seed.
    """
    if seed == 0:
        seed = random.randint(1, 2_147_483_647)
    cpid = _submit_txt2img(prompt, width, height, steps, seed)
    job_id = _register_job(cpid, prompt, "txt2img")
    return _wait_for_job(job_id)


@mcp.tool()
def img2img(
    prompt: str,
    images: list[str],
    strength: float = 0.75,
    steps: int = 4,
    seed: int = 0,
) -> list:
    """Apply a prompt to one or more existing images. Waits for all results and returns them inline.

    Each element of `images` can be:
    - URL: https://example.com/photo.jpg
    - Base64 data URI: data:image/png;base64,...
    - Job reference: job:<job_id>  (chains from a previous generate_image or img2img output)

    strength: how much to change the image. 0.1 = subtle variation, 0.75 = significant
    change, 1.0 = treat as txt2img (ignore original content). Default 0.75.

    All images are submitted at once and results returned inline when all are done.
    seed=0 picks a different random seed for each image.
    """
    if not images:
        return ["No images provided."]
    base_seed = seed if seed != 0 else random.randint(1, 2_147_483_647)
    job_ids = []
    out = []
    for i, src in enumerate(images):
        try:
            img_bytes = _resolve_image(src)
        except Exception as e:
            out.append(f"Image {i + 1}: ERROR resolving source — {e}")
            continue
        try:
            cpid = _submit_img2img(prompt, img_bytes, strength, steps, base_seed + i)
        except Exception as e:
            out.append(f"Image {i + 1}: ERROR submitting — {e}")
            continue
        job_id = _register_job(cpid, prompt, "img2img")
        job_ids.append((i + 1, job_id))

    # Wait for all jobs concurrently (poll shared registry)
    deadline = time.time() + GEN_DEADLINE
    pending = set(jid for _, jid in job_ids)
    while pending and time.time() < deadline:
        time.sleep(POLL_INTERVAL)
        with _lock:
            snapshot = {jid: dict(_jobs[jid]) for jid in pending if jid in _jobs}
        pending = {jid for jid, j in snapshot.items() if j["status"] == "running"}

    for idx, job_id in job_ids:
        with _lock:
            job = dict(_jobs.get(job_id, {}))
        elapsed = int(time.time() - job.get("submitted_at", time.time()))
        if job.get("status") == "done":
            out.append(f"Image {idx} **{job_id}** — done in {elapsed}s")
            out.append(Image(data=job["result_png"], format="png"))
        elif job.get("status") == "error":
            out.append(f"Image {idx} **{job_id}** — failed: {job.get('error','unknown')}")
        else:
            out.append(f"Image {idx} **{job_id}** — still running after {GEN_DEADLINE}s. Call get_job_status([\"{job_id}\"])")
    return out


@mcp.tool()
def get_job_status(job_ids: list[str]) -> list:
    """Get the status of one or more jobs. Completed jobs are returned with the image inline.

    For running jobs, call again in ~30 seconds.
    Completed images can be referenced in img2img via 'job:<job_id>'.
    """
    with _lock:
        snapshot = {jid: dict(_jobs[jid]) for jid in job_ids if jid in _jobs}
    out = []
    for jid in job_ids:
        job = snapshot.get(jid)
        if not job:
            out.append(f"Job **{jid}**: not found.")
            continue
        elapsed = int(time.time() - job["submitted_at"])
        if job["status"] == "done":
            url_line = f" — S3: {job['s3_url']}" if job.get("s3_url") else ""
            out.append(f"Job **{jid}**: done in {elapsed}s — {job['kind']} — {job['prompt'][:80]}{url_line}")
            out.append(Image(data=job["result_png"], format="png"))
        elif job["status"] == "error":
            out.append(f"Job **{jid}**: ERROR ({elapsed}s) — {job.get('error', 'unknown')}")
        elif job["status"] == "cancelled":
            out.append(f"Job **{jid}**: cancelled ({elapsed}s)")
        else:
            out.append(f"Job **{jid}**: running ({elapsed}s) — {job['prompt'][:60]}")
    return out


@mcp.tool()
def list_queue(limit: int = 10) -> str:
    """List the most recent jobs (all statuses) in the in-process registry.

    Shows job_id, kind, status, elapsed time, and prompt snippet.
    Use job_ids from this list with get_job_status() or as 'job:<id>' in img2img.
    """
    with _lock:
        jobs = sorted(_jobs.items(), key=lambda x: x[1]["submitted_at"], reverse=True)[:limit]
    if not jobs:
        return "No jobs in registry."
    lines = ["Recent jobs (newest first):\n"]
    for jid, j in jobs:
        elapsed = int(time.time() - j["submitted_at"])
        lines.append(f"- **{jid}** [{j['status']:9s}] {j['kind']:7s} {elapsed:5d}s  {j['prompt'][:60]}")
    return "\n".join(lines)


@mcp.tool()
def cancel_job(job_id: str) -> str:
    """Cancel a queued or running job.

    Only call this after the user has explicitly confirmed. Cannot be undone.
    """
    with _lock:
        job = _jobs.get(job_id)
    if not job:
        return f"Job {job_id}: not found."
    if job["status"] not in ("running",):
        return f"Job {job_id}: already {job['status']} — nothing to cancel."
    cpid = job["comfy_pid"]
    actions = []
    errors = []
    with httpx.Client(timeout=10) as c:
        try:
            r = c.request("DELETE", f"{COMFY}/queue", json={"delete": [cpid]})
            r.raise_for_status()
            actions.append("dequeued")
        except Exception as e:
            errors.append(f"queue delete: {e}")
        try:
            r = c.post(f"{COMFY}/interrupt")
            r.raise_for_status()
            actions.append("interrupt sent")
        except Exception as e:
            errors.append(f"interrupt: {e}")
    with _lock:
        _jobs[job_id]["status"] = "cancelled"
    msg = f"Job {job_id}: cancelled"
    if actions:
        msg += f" ({', '.join(actions)})"
    if errors:
        msg += f". Warnings: {', '.join(errors)}"
    return msg


@mcp.tool()
def list_recent_images(n: int = 5) -> list:
    """Return the n most recently completed images from the registry.

    Use this to review previous outputs or pick one as input for img2img
    by passing 'job:<job_id>' in the images list.
    """
    with _lock:
        done = [(jid, j) for jid, j in _jobs.items() if j["status"] == "done" and j["result_png"]]
    done.sort(key=lambda x: x[1]["submitted_at"], reverse=True)
    done = done[:n]
    if not done:
        return ["No completed images in the current registry. Submit a job with generate_image() first."]
    out = [f"**{len(done)} recent completed image(s)** — chain with img2img via 'job:<id>':\n"]
    for jid, j in done:
        elapsed = int(time.time() - j["submitted_at"])
        url_line = f" — S3: {j['s3_url']}" if j.get("s3_url") else ""
        out.append(f"**{jid}** ({j['kind']}, {elapsed}s ago): {j['prompt'][:80]}{url_line}")
        out.append(Image(data=j["result_png"], format="png"))
    return out


if __name__ == "__main__":
    mcp.run(transport="streamable-http")