chore: initial commit
This commit is contained in:
@@ -0,0 +1,430 @@
|
||||
"""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"),
|
||||
).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")
|
||||
Reference in New Issue
Block a user