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:
Patrick Plate
2026-04-06 10:43:36 +02:00
parent 0ff3f20589
commit c662a5237b
3 changed files with 222 additions and 43 deletions
+21
View File
@@ -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
View File
@@ -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 (110). 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 (110). 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".
"""