Content
# Universe Imagine MCP
MCP server for AI image generation. Wraps local ComfyUI as a Streamable HTTP MCP server. Works with any MCP-compatible bot client.
## Features
- **Streamable HTTP** — MCP `2024-11-05` protocol, Bearer auth
- **Synchronous generation** — `generate_image` blocks and returns the image URL directly
- **Background queue** — multiple requests accepted, processed one at a time
- **Idempotency keys** — safe retries without duplicates
- **Player-aware** — reads `player_id` from MCP initialize handshake
- **Generation history** — SQLite, with `since` filter for incremental queries
- **Health endpoint** — reports backend status, queue depth, uptime
- **Image serving** — built-in static file server for generated images
- **Idle watchdog** — stops ComfyUI after 5 minutes of inactivity to free VRAM
- **Warm start** — ComfyUI starts in background when MCP server boots
- **Discord webhook** — optional generation notifications
## Prerequisites
- Python 3.11+
- [ComfyUI](https://github.com/comfyanonymous/ComfyUI) installed locally
- At least one model downloaded (e.g. Flux, SDXL, SD3, etc.)
## Setup
```bash
# 1. Clone the repo
git clone https://github.com/BAWES-Universe/universe-imagine-mcp.git
cd universe-imagine-mcp
# 2. Create virtual environment
python3 -m venv .venv
source .venv/bin/activate
# 3. Install dependencies
pip install -r requirements.txt
# 4. Configure
cp .env.example .env
# Edit .env — at minimum set DESIGNER_MCP_TOKEN and PUBLIC_URL_BASE
# Set MODEL_UNET, MODEL_CLIP, MODEL_VAE to match your ComfyUI models
# 5. Run
python3 designer_mcp.py
```
## Configuration
All config goes in `.env`:
| Variable | Default | Description |
|----------|---------|-------------|
| `DESIGNER_MCP_TOKEN` | — | Bearer token for MCP auth (required) |
| `PUBLIC_URL_BASE` | — | Public base URL for image serving (required) |
| `HOST` | `0.0.0.0` | Listen address |
| `PORT` | `8003` | MCP server port |
| `COMFYUI_URL` | `http://127.0.0.1:8188` | ComfyUI API URL |
| `COMFYUI_DIR` | `~/ComfyUI` | ComfyUI install directory |
| `DB_PATH` | `./imagine.db` | SQLite database path |
| `OUTPUT_DIR` | `~/ComfyUI/output` | ComfyUI output directory |
| `DISCORD_WEBHOOK_URL` | — | Discord webhook for generation notifications (optional) |
| `MODEL_UNET` | `flux-2-klein-4b-fp8.safetensors` | UNet model filename |
| `MODEL_CLIP` | `Qwen3-4B-Q4_K_M.gguf` | CLIP model filename |
| `MODEL_VAE` | `flux2_klein_vae.safetensors` | VAE model filename |
## Tools
### `generate_image(prompt, width=512, height=512, idempotency_key="")` → JSON
Generate an image from a text prompt. Blocks until complete (typically 3-15 seconds) and returns the image URL directly. The image is auto-sent to the conversation by the bot.
Use `idempotency_key` to safely retry without duplicates — same key = cached result.
### `get_generation(generation_id)` → JSON
Retrieve a past generation's full record by ID.
### `list_generations(limit=10, since="")` → JSON array
List recent generations for the calling player. Optionally filter by ISO 8601 timestamp.
### `get_generation_status(generation_id)` → JSON (legacy)
Legacy polling tool kept for backwards compatibility. The current `generate_image` is synchronous and returns the URL directly — polling is not needed.
## DNS Rebinding — Monkey Patch
The server includes a monkey-patch for FastMCP's built-in DNS rebinding check (see `designer_mcp.py` lines 23-33). This is needed when running behind a reverse proxy or Cloudflare tunnel where the incoming `Host` header (e.g., `imagine.yourdomain.com`) doesn't match `localhost`.
**If you're running behind a reverse proxy:** The patch is safe because auth is handled by the `BearerAuthMiddleware`. Keep it.
**If you're running on localhost only (no tunnel/proxy):** Remove the patch entirely — delete lines 23-33 from `designer_mcp.py`.
The patch:
```python
from mcp.server.transport_security import TransportSecurityMiddleware
_orig_validate = TransportSecurityMiddleware.validate_request
async def _patched_validate(self, request, is_post=False):
return None # Skip Host header validation
TransportSecurityMiddleware.validate_request = _patched_validate
```
## Connecting to a Universe Bot
Register the MCP server on a bot via the admin API:
```json
POST /api/bots/{botId}/mcp-servers
{
"serverUrl": "https://your-domain.com/mcp",
"transport": "streamable-http",
"authType": "bearer",
"authConfig": "your-token"
}
```
Player ID is sent automatically in the MCP initialize handshake — no extra config needed.
## Architecture
```
┌─────────────┐ Streamable HTTP ┌──────────────────┐
│ Bot AI │ ◄─────────────────────── │ Imagine MCP │
│ (MCP client)│ /mcp (port 8003) │ (FastMCP) │
└─────────────┘ │ │
│ ┌────────────┐ │
┌─────────────┐ Static files │ │ Queue (1) │──┼──► _run_generation()
│ Browser │ ◄─────────────────────── │ └────────────┘ │
│ /images/* │ /images (port 8004) │ │ │
└─────────────┘ │ ▼ │
│ ┌────────────┐ │
curl /health │ │ ComfyUI │ │
─────────► /health (port 8005) │ │ (port 8188)│ │
│ └────────────┘ │
└──────────────────┘
```
## License
MIT
