feat(mcp-image-gen): add ComfyUI auto-start health check + systemd service
Option A: Add lifespan context manager to server.py - _ping_comfyui(): async health check against /system_stats - check_and_start_comfyui(): ping on startup; if down, launches ComfyUI via subprocess.Popen from COMFYUI_DIR (.venv/bin/python main.py) with HSA_OVERRIDE_GFX_VERSION=11.0.0 injected for AMD ROCm - Polls up to 30s for readiness after auto-start - New env var: COMFYUI_DIR (default ~/ComfyUI) - FastMCP lifespan= wired in; 34/34 tests still passing Option B: Add comfyui.service systemd user service file - Install: cp mcp/mcp-image-gen/comfyui.service ~/.config/systemd/user/ - Enable: systemctl --user enable --now comfyui - Sets HSA_OVERRIDE_GFX_VERSION=11.0.0, WorkingDirectory=%h/ComfyUI - Restart=on-failure, logs via journald docs: Update mcp-image-gen-ComfyUI-Setup.md - New Step 4: systemd service install + linger instructions - Step 5: manual start (moved from old Step 4) - Step 6/7 renumbered; COMFYUI_DIR env var documented - Architecture diagram added; troubleshooting rows updated
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
[Unit]
|
||||
Description=ComfyUI — Local AI Image Generation (AMD ROCm / FLUX.1-schnell)
|
||||
Documentation=https://github.com/comfyanonymous/ComfyUI
|
||||
After=network.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
WorkingDirectory=%h/ComfyUI
|
||||
ExecStart=%h/ComfyUI/.venv/bin/python main.py --listen --port 8188
|
||||
Restart=on-failure
|
||||
RestartSec=10
|
||||
|
||||
# AMD RX 7900 XTX ROCm GFX override — required for correct GPU detection
|
||||
Environment=HSA_OVERRIDE_GFX_VERSION=11.0.0
|
||||
|
||||
# Redirect output — follow with: journalctl --user -u comfyui -f
|
||||
StandardOutput=journal
|
||||
StandardError=journal
|
||||
|
||||
[Install]
|
||||
WantedBy=default.target
|
||||
+120
-33
@@ -4,16 +4,23 @@ import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import subprocess
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
import httpx
|
||||
from fastmcp import FastMCP
|
||||
from mcp.types import ImageContent, TextContent
|
||||
from pydantic import Field
|
||||
|
||||
logger = logging.getLogger("mcp-image-gen")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
@@ -23,13 +30,112 @@ COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://localhost:8188").rstrip("/")
|
||||
IMAGE_OUTPUT_DIR = os.environ.get("IMAGE_OUTPUT_DIR", "~/Pictures/mcp-generated")
|
||||
COMFYUI_TIMEOUT = int(os.environ.get("COMFYUI_TIMEOUT", "120"))
|
||||
|
||||
# Directory where ComfyUI is installed (used for auto-start only)
|
||||
# Override via COMFYUI_DIR env var. Systemd service sets this automatically.
|
||||
COMFYUI_DIR = Path(
|
||||
os.environ.get("COMFYUI_DIR", "~/ComfyUI")
|
||||
).expanduser().resolve()
|
||||
|
||||
# Maximum number of images allowed in a single batch call
|
||||
MAX_COUNT = 10
|
||||
|
||||
# Path to the bundled FLUX.1-schnell workflow template
|
||||
_WORKFLOW_PATH = Path(__file__).parent / "workflows" / "flux_schnell.json"
|
||||
|
||||
mcp = FastMCP("mcp-image-gen")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ComfyUI health check + auto-start
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def _ping_comfyui(url: str, timeout: float = 5.0) -> bool:
|
||||
"""Return True if ComfyUI is reachable at *url*/system_stats."""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
resp = await client.get(f"{url}/system_stats")
|
||||
return resp.status_code == 200
|
||||
except (httpx.ConnectError, httpx.TimeoutException, OSError):
|
||||
return False
|
||||
|
||||
|
||||
async def check_and_start_comfyui() -> None:
|
||||
"""Ping ComfyUI; if not reachable, attempt to launch it as a subprocess.
|
||||
|
||||
Called once at server startup from the lifespan context manager.
|
||||
Uses COMFYUI_DIR to locate the installation and its venv Python.
|
||||
The HSA_OVERRIDE_GFX_VERSION=11.0.0 env var is injected automatically
|
||||
for AMD ROCm / RX 7900 XTX compatibility.
|
||||
"""
|
||||
if await _ping_comfyui(COMFYUI_URL):
|
||||
logger.info("ComfyUI is already running at %s ✓", COMFYUI_URL)
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
"ComfyUI not reachable at %s — attempting to start from %s",
|
||||
COMFYUI_URL, COMFYUI_DIR,
|
||||
)
|
||||
|
||||
python = COMFYUI_DIR / ".venv" / "bin" / "python"
|
||||
main_py = COMFYUI_DIR / "main.py"
|
||||
|
||||
if not python.exists():
|
||||
logger.error(
|
||||
"ComfyUI venv Python not found at %s. "
|
||||
"Install ComfyUI first (see docs/wiki/pages/mcp-image-gen-ComfyUI-Setup.md).",
|
||||
python,
|
||||
)
|
||||
return
|
||||
if not main_py.exists():
|
||||
logger.error(
|
||||
"ComfyUI main.py not found at %s — is COMFYUI_DIR correct?",
|
||||
main_py,
|
||||
)
|
||||
return
|
||||
|
||||
# Build environment: inherit current env, set ROCm override for AMD RX 7900 XTX
|
||||
env = os.environ.copy()
|
||||
env.setdefault("HSA_OVERRIDE_GFX_VERSION", "11.0.0")
|
||||
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
[str(python), str(main_py), "--listen", "--port", "8188"],
|
||||
cwd=str(COMFYUI_DIR),
|
||||
env=env,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
start_new_session=True, # detach from MCP server process group
|
||||
)
|
||||
logger.info("ComfyUI launched (PID %d) — waiting for readiness…", proc.pid)
|
||||
except OSError as exc:
|
||||
logger.error("Failed to start ComfyUI subprocess: %s", exc)
|
||||
return
|
||||
|
||||
# Wait up to 30 s for ComfyUI to become ready (polls every 2 s)
|
||||
wait_limit = 30
|
||||
for attempt in range(wait_limit // 2):
|
||||
await asyncio.sleep(2)
|
||||
if await _ping_comfyui(COMFYUI_URL):
|
||||
logger.info(
|
||||
"ComfyUI ready at %s after ~%ds ✓", COMFYUI_URL, (attempt + 1) * 2
|
||||
)
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
"ComfyUI did not respond within %ds. "
|
||||
"Generation calls will fail until it is ready. "
|
||||
"Check logs: journalctl --user -u comfyui -f",
|
||||
wait_limit,
|
||||
)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app):
|
||||
"""FastMCP lifespan: run ComfyUI health check at server startup."""
|
||||
await check_and_start_comfyui()
|
||||
yield # server is live here
|
||||
# Nothing to tear down — ComfyUI is managed by systemd, not this process
|
||||
|
||||
|
||||
mcp = FastMCP("mcp-image-gen", lifespan=lifespan)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -332,40 +438,22 @@ async def _generate_single(
|
||||
|
||||
@mcp.tool()
|
||||
async def generate_image(
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
steps: int = 4,
|
||||
model: str = "flux1-schnell.safetensors",
|
||||
seed: int = -1,
|
||||
negative_prompt: str = "",
|
||||
output_dir: str = "",
|
||||
name: str = "",
|
||||
count: int = 1,
|
||||
prompt: Annotated[str, Field(description="Text description of the image to generate.")],
|
||||
width: Annotated[int, Field(description="Image width in pixels (default: 1024).")] = 1024,
|
||||
height: Annotated[int, Field(description="Image height in pixels (default: 1024).")] = 1024,
|
||||
steps: Annotated[int, Field(description="Number of inference steps. FLUX.1-schnell works well at 4.")] = 4,
|
||||
model: Annotated[str, Field(description="ComfyUI model filename (default: flux1-schnell.safetensors).")] = "flux1-schnell.safetensors",
|
||||
seed: Annotated[int, Field(description="Random seed for reproducibility. -1 = random. When count > 1 and seed != -1, seeds are incremented per image (seed, seed+1, seed+2, ...) to produce deterministic variation.")] = -1,
|
||||
negative_prompt: Annotated[str, Field(description="Things to exclude from the image (optional).")] = "",
|
||||
output_dir: Annotated[str, Field(description="Override output directory. Defaults to IMAGE_OUTPUT_DIR env var or ~/Pictures/mcp-generated.")] = "",
|
||||
name: Annotated[str, Field(description="Optional filename prefix. Saved as {name}_{timestamp}_{seed}.png. Useful to avoid confusion with auto-generated timestamp filenames.")] = "",
|
||||
count: Annotated[int, Field(description="Number of images to generate (1–10). Each image is generated sequentially. Partial failures are returned inline — the batch continues even if one image fails.")] = 1,
|
||||
) -> list:
|
||||
"""Generate an image from a text prompt using ComfyUI.
|
||||
|
||||
Returns both a file path (for persistence) and an inline base64 image
|
||||
(for display in Claude / Roo Code chat).
|
||||
|
||||
Args:
|
||||
prompt: Text description of the image to generate.
|
||||
width: Image width in pixels (default: 1024).
|
||||
height: Image height in pixels (default: 1024).
|
||||
steps: Number of inference steps. FLUX.1-schnell works well at 4.
|
||||
model: ComfyUI model filename (default: flux1-schnell.safetensors).
|
||||
seed: Random seed for reproducibility. -1 = random.
|
||||
When count > 1 and seed != -1, seeds are incremented per image
|
||||
(seed, seed+1, seed+2, ...) to produce deterministic variation.
|
||||
negative_prompt: Things to exclude from the image (optional).
|
||||
output_dir: Override output directory. Defaults to IMAGE_OUTPUT_DIR env var
|
||||
or ~/Pictures/mcp-generated.
|
||||
name: Optional filename prefix. Saved as {name}_{timestamp}_{seed}.png.
|
||||
Useful to avoid confusion with auto-generated timestamp filenames.
|
||||
count: Number of images to generate (1–10). Each image is generated
|
||||
sequentially. Partial failures are returned inline — the batch
|
||||
continues even if one image fails.
|
||||
|
||||
Returns:
|
||||
Flat interleaved list: [TextContent1, ImageContent1, TextContent2, ImageContent2, ...]
|
||||
On error for any single image, that slot contains only [TextContent(error)].
|
||||
@@ -442,12 +530,11 @@ async def list_available_models() -> list[str]:
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_generation_status(prompt_id: str) -> dict:
|
||||
async def get_generation_status(
|
||||
prompt_id: Annotated[str, Field(description="The prompt ID returned by a previous generate_image call.")],
|
||||
) -> dict:
|
||||
"""Check the status of a queued or running generation job.
|
||||
|
||||
Args:
|
||||
prompt_id: The prompt ID returned by a previous generate_image call.
|
||||
|
||||
Returns:
|
||||
Dict with 'status' key: "pending", "running", "completed", or "not_found".
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user