feat(mcp-image-gen): merge name and count params into main
This commit is contained in:
Patrick Plate
@@ -0,0 +1,319 @@
|
||||
# Assessment: Expand `generate_image` with `name` and `count` Parameters
|
||||
|
||||
*Author: Lumen | Date: 2026-04-06 | Ticket: —*
|
||||
*BigMind Session: `00070c37-b013-4342-a8ae-f81da0e3180d`*
|
||||
*Status: 🔵 DRAFT — awaiting Patrick review*
|
||||
|
||||
---
|
||||
|
||||
## 1. Problem Statement
|
||||
|
||||
The current [`generate_image()`](mcp/mcp-image-gen/src/server.py:133) tool generates a single image and saves it with an auto-generated filename of `{timestamp}_{seed}.png`. Two common workflows are not yet supported:
|
||||
|
||||
1. **Named outputs** — When generating thematic sets (Lumen profile images, wiki banners, concept art), the caller wants a meaningful prefix in the filename (e.g., `lumen_profile_20260406_140236_2409122067.png`) rather than a bare timestamp. This also enables grouping output by purpose in the directory listing.
|
||||
|
||||
2. **Batch generation** — Generating multiple variations of the same prompt in one tool call is a common creative workflow. Currently, the caller must invoke `generate_image` N times with separate tool calls, which is verbose and loses the semantic grouping.
|
||||
|
||||
**Goal:** Add two optional parameters — `name` (filename prefix string) and `count` (integer repetitions) — to `generate_image` with minimal disruption to existing behaviour and test coverage.
|
||||
|
||||
---
|
||||
|
||||
## 2. Requirements
|
||||
|
||||
### 2.1 Functional Requirements
|
||||
|
||||
| ID | Requirement |
|
||||
|----|-------------|
|
||||
| F-1 | `name` parameter (default `""`) prepends a sanitized label to the output filename |
|
||||
| F-2 | When `name=""` (default), filename format is unchanged: `{timestamp}_{seed}.png` |
|
||||
| F-3 | When `name="lumen_profile"`, filename format is: `lumen_profile_{timestamp}_{seed}.png` |
|
||||
| F-4 | `count` parameter (default `1`) generates N images sequentially |
|
||||
| F-5 | When `count=1` (default), return value is identical to the current `[TextContent, ImageContent]` |
|
||||
| F-6 | When `count=N > 1`, return value is a flat list: `[Text1, Image1, Text2, Image2, ..., TextN, ImageN]` |
|
||||
| F-7 | When `count>1` and `seed=-1`, each image gets an independently random seed |
|
||||
| F-8 | When `count>1` and a fixed `seed` is provided, images use `seed`, `seed+1`, `seed+2`, … to produce deterministic variation |
|
||||
| F-9 | `count` is capped at a maximum (proposed: 10) to prevent runaway generation |
|
||||
| F-10 | `name` is sanitized: non-alphanumeric characters (except `-` and `_`) are stripped/replaced; max 64 chars |
|
||||
| F-11 | Partial success: if one image in a batch fails, the error is returned as a `TextContent` error item in that position rather than aborting the whole batch |
|
||||
| F-12 | The TextContent for each image in a batch includes the 1-of-N index: `[1/3] Generated: ...` |
|
||||
|
||||
### 2.2 Non-Functional Requirements
|
||||
|
||||
| ID | Requirement |
|
||||
|----|-------------|
|
||||
| NF-1 | Sequential generation — no concurrent ComfyUI submissions (ComfyUI queues internally; parallel MCP submissions would complicate polling) |
|
||||
| NF-2 | Backward compatibility — all existing callers with no `name`/`count` args produce identical output |
|
||||
| NF-3 | All existing 19 tests must continue to pass without modification |
|
||||
| NF-4 | New tests must cover: name prefix in filename, count=2 success, count with fixed seed increments, count with partial failure, name sanitization, count cap enforcement |
|
||||
| NF-5 | MCP tool schema (visible in Claude/Roo Code) must surface clear descriptions for the new params |
|
||||
|
||||
---
|
||||
|
||||
## 3. Affected Files
|
||||
|
||||
| File | Change Type | Description |
|
||||
|------|-------------|-------------|
|
||||
| [`mcp/mcp-image-gen/src/server.py`](mcp/mcp-image-gen/src/server.py:133) | Modify | Add `name: str = ""` and `count: int = 1` params to `generate_image()`; add `_sanitize_name()` helper; extract `_generate_single()` inner logic |
|
||||
| [`mcp/mcp-image-gen/tests/test_server.py`](mcp/mcp-image-gen/tests/test_server.py:1) | Modify | Add 6+ new test cases covering new parameters |
|
||||
| [`mcp/mcp-image-gen/README.md`](mcp/mcp-image-gen/README.md) | Modify | Update `generate_image` tool documentation table |
|
||||
| [`docs/wiki/pages/mcp-image-gen.md`](docs/wiki/pages/mcp-image-gen.md) | Modify | Update tool reference table with new parameters |
|
||||
|
||||
No schema changes, no new dependencies, no workflow JSON changes.
|
||||
|
||||
---
|
||||
|
||||
## 4. Design Decisions
|
||||
|
||||
### 4.1 Filename Convention with `name`
|
||||
|
||||
**Current:** `{timestamp}_{seed}.png`
|
||||
**Proposed:** `{sanitized_name}_{timestamp}_{seed}.png` (when `name` is provided)
|
||||
|
||||
The `name` is placed as a **prefix** rather than suffix so directory `ls` output groups named sets together alphabetically:
|
||||
```
|
||||
lumen_profile_20260406_140236_2409122067.png
|
||||
lumen_profile_20260406_140258_764633840.png
|
||||
wiki_banner_20260406_141000_1234567.png
|
||||
```
|
||||
|
||||
**Sanitization rule:** `re.sub(r'[^a-zA-Z0-9_-]', '_', name)[:64]` — replaces any character that is not alphanumeric, dash, or underscore with `_`, then truncates to 64 chars.
|
||||
|
||||
### 4.2 Seed Behaviour for Batch Generation
|
||||
|
||||
| Scenario | Behaviour |
|
||||
|----------|-----------|
|
||||
| `count=3, seed=-1` | Each call to `build_flux_workflow` gets `seed=-1` → 3 independent random seeds |
|
||||
| `count=3, seed=42` | Seeds are 42, 43, 44 — deterministic, reproducible variation |
|
||||
|
||||
This follows the convention of most image generation tools (e.g., ComfyUI's own batch seed increment).
|
||||
|
||||
### 4.3 Return Structure for `count > 1`
|
||||
|
||||
Return a **flat interleaved list**: `[Text1, Image1, Text2, Image2]`
|
||||
|
||||
**Rationale:** MCP content lists are flat arrays. Claude/Roo Code renders them sequentially — a flat list means each image appears immediately below its metadata line. A nested structure would require the caller to unwrap it.
|
||||
|
||||
**For `count=1` (default):** Behaviour is identical to today — `[TextContent, ImageContent]`. No caller breakage.
|
||||
|
||||
### 4.4 Refactoring: Extract `_generate_single()`
|
||||
|
||||
The current `generate_image` function is 180+ lines of inline logic. To support `count`, the inner pipeline (queue → poll → history → download → save → encode) will be extracted to a private `async def _generate_single(prompt, ..., index, total)` coroutine. `generate_image` then loops `count` times calling `_generate_single` and accumulates results.
|
||||
|
||||
This refactoring:
|
||||
- Makes the count loop clean (`results.extend(await _generate_single(...))`)
|
||||
- Makes partial failure handling straightforward (catch per iteration)
|
||||
- Improves testability of the single-image path
|
||||
|
||||
### 4.5 Maximum Count Cap
|
||||
|
||||
Cap `count` at **10**. Rationale:
|
||||
- FLUX.1-schnell takes ~10–35s per image on RX 7900 XTX → 10 images ≈ 100–350s maximum
|
||||
- MCP tool call timeout in Roo Code defaults to 5 minutes — 10 images is safe margin
|
||||
- ComfyUI queues them internally; the MCP server polls sequentially, not in parallel
|
||||
|
||||
When `count > 10`, the tool returns a single `TextContent` error immediately (no images generated) with message: `"count={N} exceeds maximum of 10. Reduce count and retry."`
|
||||
|
||||
---
|
||||
|
||||
## 5. Implementation Plan
|
||||
|
||||
### Step 1 — Add `_sanitize_name()` helper
|
||||
|
||||
```python
|
||||
import re
|
||||
|
||||
def _sanitize_name(name: str) -> str:
|
||||
"""Sanitize a name for use as a filename prefix."""
|
||||
sanitized = re.sub(r'[^a-zA-Z0-9_-]', '_', name)
|
||||
return sanitized[:64]
|
||||
```
|
||||
|
||||
Location: [`server.py`](mcp/mcp-image-gen/src/server.py:95), after `build_flux_workflow()` (pure function section).
|
||||
|
||||
### Step 2 — Extract `_generate_single()` coroutine
|
||||
|
||||
Extract the body of the current `generate_image` (lines 162–310) into:
|
||||
|
||||
```python
|
||||
async def _generate_single(
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
steps: int,
|
||||
model: str,
|
||||
seed: int,
|
||||
negative_prompt: str,
|
||||
resolved_output_dir: Path,
|
||||
filename_prefix: str,
|
||||
index: int,
|
||||
total: int,
|
||||
) -> list:
|
||||
```
|
||||
|
||||
The `filename` construction changes to:
|
||||
```python
|
||||
filename = f"{filename_prefix}{timestamp}_{actual_seed}.png"
|
||||
# where filename_prefix = f"{sanitized_name}_" if sanitized_name else ""
|
||||
```
|
||||
|
||||
The `TextContent` text changes when `total > 1`:
|
||||
```python
|
||||
prefix_label = f"[{index}/{total}] " if total > 1 else ""
|
||||
text = f"{prefix_label}Generated: {out_path}\nSeed: ..."
|
||||
```
|
||||
|
||||
### Step 3 — Update `generate_image()` signature
|
||||
|
||||
```python
|
||||
@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,
|
||||
) -> list:
|
||||
```
|
||||
|
||||
Body of `generate_image` becomes:
|
||||
|
||||
```python
|
||||
# Validate count
|
||||
MAX_COUNT = 10
|
||||
if count < 1 or count > MAX_COUNT:
|
||||
return [TextContent(type="text", text=f"count={count} is invalid. Must be 1–{MAX_COUNT}.")]
|
||||
|
||||
sanitized_name = _sanitize_name(name) if name else ""
|
||||
filename_prefix = f"{sanitized_name}_" if sanitized_name else ""
|
||||
resolved_output_dir = Path(output_dir or IMAGE_OUTPUT_DIR).expanduser().resolve()
|
||||
|
||||
results = []
|
||||
for i in range(1, count + 1):
|
||||
actual_seed = seed if seed == -1 else seed + (i - 1)
|
||||
items = await _generate_single(
|
||||
prompt=prompt, width=width, height=height, steps=steps,
|
||||
model=model, seed=actual_seed, negative_prompt=negative_prompt,
|
||||
resolved_output_dir=resolved_output_dir,
|
||||
filename_prefix=filename_prefix, index=i, total=count,
|
||||
)
|
||||
results.extend(items)
|
||||
|
||||
return results
|
||||
```
|
||||
|
||||
### Step 4 — Write new tests
|
||||
|
||||
Add to [`test_server.py`](mcp/mcp-image-gen/tests/test_server.py:550):
|
||||
|
||||
| Test | Description |
|
||||
|------|-------------|
|
||||
| `test_generate_image_with_name` | `name="lumen"` → filename starts with `lumen_` |
|
||||
| `test_generate_image_name_sanitization` | `name="my image! v2"` → `my_image__v2_` prefix |
|
||||
| `test_generate_image_count_2_success` | `count=2` → 4 items in result, 2 files saved |
|
||||
| `test_generate_image_count_fixed_seed` | `count=2, seed=42` → seeds 42 and 43 in filenames |
|
||||
| `test_generate_image_count_partial_failure` | `count=2`, second POST fails → 2 items (success) + 1 item (error) |
|
||||
| `test_generate_image_count_cap_exceeded` | `count=11` → single TextContent error, no generation |
|
||||
| `test_generate_image_count_0_invalid` | `count=0` → single TextContent error |
|
||||
| `test_generate_image_name_and_count_combined` | `name="banner", count=2` → both files prefixed `banner_` |
|
||||
|
||||
### Step 5 — Update documentation
|
||||
|
||||
- Update `generate_image` docstring in [`server.py`](mcp/mcp-image-gen/src/server.py:144) to document `name` and `count`
|
||||
- Update parameter table in [`README.md`](mcp/mcp-image-gen/README.md)
|
||||
- Update tool reference in [`docs/wiki/pages/mcp-image-gen.md`](docs/wiki/pages/mcp-image-gen.md)
|
||||
|
||||
### Step 6 — Run full test suite
|
||||
|
||||
```bash
|
||||
cd mcp/mcp-image-gen && uv run pytest tests/ -v --tb=short
|
||||
```
|
||||
|
||||
All 19 existing + 8 new = **27 tests** must pass.
|
||||
|
||||
### Step 7 — Commit and push
|
||||
|
||||
Branch: `feat/mcp-image-gen/generate-image-name-count`
|
||||
Commit: `feat(mcp-image-gen): add name and count params to generate_image`
|
||||
|
||||
---
|
||||
|
||||
## 6. Risks
|
||||
|
||||
| Risk | Likelihood | Impact | Mitigation |
|
||||
|------|------------|--------|------------|
|
||||
| Partial batch failure leaves orphaned files on disk | Medium | Low | Files for successful images are kept; error TextContent clearly identifies which index failed. No cleanup needed — partial results are useful. |
|
||||
| `count` loop adds significant latency visible in Roo Code | Medium | Medium | Document expected time: `count × ~15s`. MCP timeout is 5 min; max 10 images ≈ 150s. Still within limit. |
|
||||
| Seed increment wraps around at `2^32` | Very Low | Low | `(seed + i - 1) % 2**32` — add modulo guard in `_generate_single` |
|
||||
| `_generate_single` refactor introduces regression in existing tests | Low | High | Existing test fixtures mock ComfyUI endpoints — as long as the HTTP call sequence is unchanged, respx mocks will match. Verify each existing test still passes before adding new ones. |
|
||||
| `name` with only special chars becomes empty after sanitization | Low | Medium | After sanitization, if result is empty string, treat as unnamed (no prefix). Add assertion in `_sanitize_name` to return `""` for all-whitespace/special inputs. |
|
||||
| MCP tool schema change breaks existing callers | Very Low | Low | New params are optional with defaults — backward compatible. Roo Code re-reads schema on server restart. |
|
||||
|
||||
---
|
||||
|
||||
## 7. Alternatives Considered
|
||||
|
||||
### 7.1 Separate `generate_images_batch()` Tool (Rejected)
|
||||
|
||||
Add a new tool instead of expanding `generate_image`.
|
||||
|
||||
**Pros:** Clean separation, no refactoring of existing tool.
|
||||
**Cons:** Two tools for the same backend; callers must learn two tool names; MCP tool list grows. The MCP convention favours extending existing tools with optional parameters rather than proliferating tools.
|
||||
|
||||
**Verdict:** Rejected. Optional parameters with backward-compatible defaults is the right pattern here.
|
||||
|
||||
### 7.2 Return Grouped List of Lists for `count > 1` (Rejected)
|
||||
|
||||
Return `[[Text1, Image1], [Text2, Image2]]` for batch results.
|
||||
|
||||
**Pros:** Caller can index by image number cleanly.
|
||||
**Cons:** MCP content type is a flat `list[ContentBlock]`. FastMCP does not support nested lists in tool returns — they would be serialized as strings, not rendered. Roo Code renders content sequentially; flat interleaved is the idiomatic structure.
|
||||
|
||||
**Verdict:** Rejected. Flat interleaved list `[Text1, Image1, Text2, Image2]` is MCP-idiomatic.
|
||||
|
||||
### 7.3 Parallel ComfyUI Submission for Batch (Rejected)
|
||||
|
||||
Submit all `count` prompts to ComfyUI simultaneously (async tasks), then collect results in order.
|
||||
|
||||
**Pros:** Faster if ComfyUI supports parallel queue processing (it does).
|
||||
**Cons:** ComfyUI processes one job at a time on a single GPU regardless — parallel submission just fills the queue. Polling becomes complex (N polling loops). Error handling harder. Out-of-order completions break index alignment.
|
||||
|
||||
**Verdict:** Rejected for v1. Sequential submission is simpler, correct, and produces no worse throughput. Can revisit if ComfyUI gains true parallel processing support.
|
||||
|
||||
### 7.4 Name as Subdirectory Instead of Filename Prefix (Rejected)
|
||||
|
||||
When `name="lumen"`, save to `output_dir/lumen/` instead of `output_dir/lumen_*.png`.
|
||||
|
||||
**Pros:** Better directory organisation for large sets.
|
||||
**Cons:** Complicates the implementation (directory creation per name), changes the return path format, breaks callers who assume a flat output directory. Adds complexity for minimal gain at `count ≤ 10`.
|
||||
|
||||
**Verdict:** Rejected for v1. Prefix approach is simpler and equally readable.
|
||||
|
||||
---
|
||||
|
||||
## 8. Success Criteria
|
||||
|
||||
| Criterion | Measure |
|
||||
|-----------|---------|
|
||||
| All 27 tests pass | `uv run pytest tests/ -v` exits 0 |
|
||||
| `name="lumen"` → file starts with `lumen_` | Assert in `test_generate_image_with_name` |
|
||||
| `count=2` → 4 content items, 2 files | Assert `len(result) == 4`, `len(glob("*.png")) == 2` |
|
||||
| `count=2, seed=42` → seeds 42 and 43 | Assert seed values in TextContent |
|
||||
| `count=11` → error TextContent, no ComfyUI call | Assert `len(result) == 1`, no `/api/prompt` mock hit |
|
||||
| Backward compat: existing callers unaffected | All 19 existing tests pass without modification |
|
||||
| MCP tool schema shows `name` and `count` params | Visible in Roo Code tool list after server restart |
|
||||
|
||||
---
|
||||
|
||||
## 9. Open Questions
|
||||
|
||||
| # | Question | Owner | Priority |
|
||||
|---|----------|-------|----------|
|
||||
| Q1 | Should `count=0` be an error, or silently return `[]` (empty list)? | Patrick | Low — assessment recommends error for clarity |
|
||||
| Q2 | Max count cap: 10 or higher? 10 ≈ 150s max at 15s/image — feels right, but could be raised to 20 for batch profile image sets. | Patrick | Medium |
|
||||
| Q3 | Should partial batch failure stop remaining iterations, or always complete all N? | Patrick | Medium — assessment recommends continue (partial success) |
|
||||
| Q4 | Should `name` parameter also tag the TextContent output text, e.g. `[lumen_profile 1/3] Generated: ...`? | Patrick | Low |
|
||||
+156
-44
@@ -6,6 +6,7 @@ import copy
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
@@ -22,6 +23,9 @@ 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"))
|
||||
|
||||
# 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"
|
||||
|
||||
@@ -126,46 +130,59 @@ def build_flux_workflow(
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tools
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@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 = "",
|
||||
) -> list:
|
||||
"""Generate an image from a text prompt using ComfyUI.
|
||||
def _sanitize_name(name: str) -> str:
|
||||
"""Sanitize a user-provided name for safe use in filenames.
|
||||
|
||||
Returns both a file path (for persistence) and an inline base64 image
|
||||
(for display in Claude / Roo Code chat).
|
||||
Replaces whitespace with underscores, strips any characters that are not
|
||||
alphanumeric, underscores, or hyphens, and collapses consecutive
|
||||
underscores/hyphens. Returns empty string if nothing usable remains.
|
||||
"""
|
||||
name = name.strip()
|
||||
name = re.sub(r"\s+", "_", name) # spaces → underscores
|
||||
name = re.sub(r"[^\w\-]", "", name) # strip non-alphanum/underscore/hyphen
|
||||
name = re.sub(r"[_\-]{2,}", "_", name) # collapse runs
|
||||
name = name.strip("_-") # trim leading/trailing separators
|
||||
return name[:64] # cap at 64 chars
|
||||
|
||||
|
||||
def _build_filename(name: str, timestamp: str, actual_seed: int) -> str:
|
||||
"""Build an output filename from optional name, timestamp and seed."""
|
||||
sanitized = _sanitize_name(name)
|
||||
if sanitized:
|
||||
return f"{sanitized}_{timestamp}_{actual_seed}.png"
|
||||
return f"{timestamp}_{actual_seed}.png"
|
||||
|
||||
|
||||
async def _generate_single(
|
||||
client: ComfyUIClient,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
steps: int,
|
||||
seed: int,
|
||||
model: str,
|
||||
resolved_output_dir: Path,
|
||||
name: str,
|
||||
label: str,
|
||||
) -> list:
|
||||
"""Generate a single image and return [TextContent, ImageContent] or [TextContent] on error.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
Returns:
|
||||
[TextContent(path + metadata), ImageContent(base64 PNG)]
|
||||
client: ComfyUIClient instance.
|
||||
prompt: Positive text prompt.
|
||||
negative_prompt: Negative text prompt.
|
||||
width / height: Image dimensions.
|
||||
steps: Inference steps.
|
||||
seed: Seed value (-1 = random).
|
||||
model: ComfyUI model filename.
|
||||
resolved_output_dir: Resolved output directory Path.
|
||||
name: User-supplied name prefix (unsanitized).
|
||||
label: Human-readable label for TextContent prefix (e.g. "[lumen 1/3]").
|
||||
"""
|
||||
# Resolve output directory
|
||||
resolved_output_dir = Path(
|
||||
output_dir or IMAGE_OUTPUT_DIR
|
||||
).expanduser().resolve()
|
||||
|
||||
client = ComfyUIClient(COMFYUI_URL)
|
||||
|
||||
# Build and submit workflow
|
||||
try:
|
||||
workflow = build_flux_workflow(
|
||||
@@ -178,14 +195,13 @@ async def generate_image(
|
||||
model=model,
|
||||
)
|
||||
actual_seed = workflow["_meta"]["actual_seed"]
|
||||
|
||||
prompt_id = await client.queue_prompt(workflow)
|
||||
except httpx.ConnectError:
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=(
|
||||
f"ComfyUI not reachable at {COMFYUI_URL}. "
|
||||
f"{label} ComfyUI not reachable at {COMFYUI_URL}. "
|
||||
"Start it with: python main.py --listen"
|
||||
),
|
||||
)
|
||||
@@ -194,7 +210,7 @@ async def generate_image(
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"ComfyUI returned an error: {e.response.status_code} — {e.response.text}",
|
||||
text=f"{label} ComfyUI returned an error: {e.response.status_code} — {e.response.text}",
|
||||
)
|
||||
]
|
||||
|
||||
@@ -207,7 +223,7 @@ async def generate_image(
|
||||
TextContent(
|
||||
type="text",
|
||||
text=(
|
||||
f"Generation timed out after {COMFYUI_TIMEOUT}s. "
|
||||
f"{label} Generation timed out after {COMFYUI_TIMEOUT}s. "
|
||||
f"prompt_id={prompt_id} — use get_generation_status to check"
|
||||
),
|
||||
)
|
||||
@@ -236,7 +252,7 @@ async def generate_image(
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"Failed to retrieve generation history: {e}",
|
||||
text=f"{label} Failed to retrieve generation history: {e}",
|
||||
)
|
||||
]
|
||||
|
||||
@@ -255,7 +271,7 @@ async def generate_image(
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"No output image found in history for prompt_id={prompt_id}",
|
||||
text=f"{label} No output image found in history for prompt_id={prompt_id}",
|
||||
)
|
||||
]
|
||||
|
||||
@@ -270,7 +286,7 @@ async def generate_image(
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"Failed to download generated image: {e}",
|
||||
text=f"{label} Failed to download generated image: {e}",
|
||||
)
|
||||
]
|
||||
|
||||
@@ -278,14 +294,14 @@ async def generate_image(
|
||||
try:
|
||||
resolved_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"{timestamp}_{actual_seed}.png"
|
||||
filename = _build_filename(name, timestamp, actual_seed)
|
||||
out_path = resolved_output_dir / filename
|
||||
out_path.write_bytes(image_bytes)
|
||||
except OSError as e:
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"Cannot write to output directory: {resolved_output_dir} — {e}",
|
||||
text=f"{label} Cannot write to output directory: {resolved_output_dir} — {e}",
|
||||
)
|
||||
]
|
||||
|
||||
@@ -296,7 +312,7 @@ async def generate_image(
|
||||
TextContent(
|
||||
type="text",
|
||||
text=(
|
||||
f"Generated: {out_path}\n"
|
||||
f"{label} Generated: {out_path}\n"
|
||||
f"Seed: {actual_seed}\n"
|
||||
f"Elapsed: {elapsed:.1f}s\n"
|
||||
f"Size: {width}x{height}, Steps: {steps}, Model: {model}"
|
||||
@@ -310,6 +326,102 @@ async def generate_image(
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tools
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@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,
|
||||
) -> 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)].
|
||||
"""
|
||||
# Validate count
|
||||
if count < 1:
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"count must be at least 1 (got {count}).",
|
||||
)
|
||||
]
|
||||
if count > MAX_COUNT:
|
||||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"count must be at most {MAX_COUNT} (got {count}). Use multiple calls for larger batches.",
|
||||
)
|
||||
]
|
||||
|
||||
# Resolve output directory once
|
||||
resolved_output_dir = Path(
|
||||
output_dir or IMAGE_OUTPUT_DIR
|
||||
).expanduser().resolve()
|
||||
|
||||
client = ComfyUIClient(COMFYUI_URL)
|
||||
|
||||
results = []
|
||||
for i in range(1, count + 1):
|
||||
# Compute seed for this image:
|
||||
# - seed=-1 → each image gets an independent random seed
|
||||
# - fixed seed → increment by i-1 for deterministic variation across the batch
|
||||
image_seed = seed if seed == -1 else seed + (i - 1)
|
||||
|
||||
label = f"[{_sanitize_name(name) or 'image'} {i}/{count}]" if count > 1 else (
|
||||
f"[{_sanitize_name(name)}]" if _sanitize_name(name) else ""
|
||||
)
|
||||
|
||||
single_result = await _generate_single(
|
||||
client=client,
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
steps=steps,
|
||||
seed=image_seed,
|
||||
model=model,
|
||||
resolved_output_dir=resolved_output_dir,
|
||||
name=name,
|
||||
label=label,
|
||||
)
|
||||
results.extend(single_result)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def list_available_models() -> list[str]:
|
||||
"""List all checkpoint models available in ComfyUI.
|
||||
|
||||
@@ -14,6 +14,8 @@ import respx
|
||||
import server
|
||||
from server import (
|
||||
ComfyUIClient,
|
||||
_build_filename,
|
||||
_sanitize_name,
|
||||
build_flux_workflow,
|
||||
generate_image,
|
||||
get_generation_status,
|
||||
@@ -100,6 +102,74 @@ def test_random_seed_generated():
|
||||
assert "_meta" in wf2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _sanitize_name — pure function
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_sanitize_name_basic():
|
||||
"""Simple alphanumeric name passes through unchanged."""
|
||||
assert _sanitize_name("lumen_profile") == "lumen_profile"
|
||||
|
||||
|
||||
def test_sanitize_name_spaces_to_underscores():
|
||||
"""Spaces are converted to underscores."""
|
||||
assert _sanitize_name("my cool image") == "my_cool_image"
|
||||
|
||||
|
||||
def test_sanitize_name_special_chars_stripped():
|
||||
"""Special characters (!, @, #, etc.) are stripped."""
|
||||
result = _sanitize_name("hello! world@2024#")
|
||||
assert "!" not in result
|
||||
assert "@" not in result
|
||||
assert "#" not in result
|
||||
assert "hello" in result
|
||||
assert "world" in result
|
||||
|
||||
|
||||
def test_sanitize_name_empty_returns_empty():
|
||||
"""Empty string or whitespace-only returns empty string."""
|
||||
assert _sanitize_name("") == ""
|
||||
assert _sanitize_name(" ") == ""
|
||||
|
||||
|
||||
def test_sanitize_name_collapse_underscores():
|
||||
"""Multiple consecutive underscores/hyphens are collapsed to one."""
|
||||
result = _sanitize_name("lumen__profile")
|
||||
assert "__" not in result
|
||||
|
||||
|
||||
def test_sanitize_name_truncates_at_64():
|
||||
"""Names longer than 64 chars are truncated."""
|
||||
long_name = "a" * 100
|
||||
result = _sanitize_name(long_name)
|
||||
assert len(result) <= 64
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _build_filename — pure function
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_build_filename_with_name():
|
||||
"""When name is provided, filename includes it as prefix."""
|
||||
filename = _build_filename("lumen", "20260406_120000", 12345)
|
||||
assert filename == "lumen_20260406_120000_12345.png"
|
||||
|
||||
|
||||
def test_build_filename_without_name():
|
||||
"""When name is empty, filename is timestamp_seed.png."""
|
||||
filename = _build_filename("", "20260406_120000", 12345)
|
||||
assert filename == "20260406_120000_12345.png"
|
||||
assert not filename.startswith("_")
|
||||
|
||||
|
||||
def test_build_filename_sanitizes_name():
|
||||
"""Name with spaces and special chars is sanitized before use in filename."""
|
||||
filename = _build_filename("my image!", "20260406_120000", 99)
|
||||
assert "!" not in filename
|
||||
assert "my_image" in filename
|
||||
assert filename.endswith("_20260406_120000_99.png")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# list_available_models
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -211,7 +281,255 @@ def test_get_output_directory_custom(monkeypatch, tmp_path):
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# generate_image
|
||||
# generate_image — count/name validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_count_zero_returns_error():
|
||||
"""count=0 → returns error TextContent without calling ComfyUI."""
|
||||
result = await generate_image(prompt="a cat", count=0)
|
||||
assert len(result) == 1
|
||||
assert "count must be at least 1" in result[0].text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_count_exceeds_max_returns_error():
|
||||
"""count=11 (> MAX_COUNT=10) → returns error TextContent without calling ComfyUI."""
|
||||
result = await generate_image(prompt="a cat", count=11)
|
||||
assert len(result) == 1
|
||||
assert "at most 10" in result[0].text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# generate_image — name parameter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@respx.mock
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_with_name(
|
||||
tmp_path, sample_image_bytes, mock_history_response, queue_empty, monkeypatch
|
||||
):
|
||||
"""name param → saved file has name as prefix."""
|
||||
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path))
|
||||
|
||||
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
|
||||
return_value=httpx.Response(200, json={"prompt_id": "name-test-uuid"})
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
|
||||
return_value=httpx.Response(200, json=queue_empty)
|
||||
)
|
||||
mock_history_named = {
|
||||
"name-test-uuid": {
|
||||
"outputs": {
|
||||
"9": {
|
||||
"images": [
|
||||
{
|
||||
"filename": "mcp-image-gen_00001_.png",
|
||||
"subfolder": "",
|
||||
"type": "output",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"status": {"completed": True},
|
||||
}
|
||||
}
|
||||
respx.get(f"{COMFYUI_BASE}/api/history/name-test-uuid").mock(
|
||||
return_value=httpx.Response(200, json=mock_history_named)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/view").mock(
|
||||
return_value=httpx.Response(200, content=sample_image_bytes)
|
||||
)
|
||||
|
||||
result = await generate_image(
|
||||
prompt="lumen portrait",
|
||||
name="lumen_profile",
|
||||
output_dir=str(tmp_path),
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
saved_files = list(tmp_path.glob("lumen_profile_*.png"))
|
||||
assert len(saved_files) == 1, f"Expected 1 file with 'lumen_profile_' prefix, got: {list(tmp_path.glob('*.png'))}"
|
||||
# Path in TextContent also has the name prefix
|
||||
assert "lumen_profile_" in result[0].text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# generate_image — count=2 batch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@respx.mock
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_count_2(
|
||||
tmp_path, sample_image_bytes, queue_empty, monkeypatch
|
||||
):
|
||||
"""count=2 → returns 4 content items (Text+Image per image), 2 files saved."""
|
||||
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path))
|
||||
|
||||
# First image
|
||||
mock_history_1 = {
|
||||
"uuid-batch-1": {
|
||||
"outputs": {"9": {"images": [{"filename": "img1.png", "subfolder": "", "type": "output"}]}},
|
||||
"status": {"completed": True},
|
||||
}
|
||||
}
|
||||
# Second image
|
||||
mock_history_2 = {
|
||||
"uuid-batch-2": {
|
||||
"outputs": {"9": {"images": [{"filename": "img2.png", "subfolder": "", "type": "output"}]}},
|
||||
"status": {"completed": True},
|
||||
}
|
||||
}
|
||||
|
||||
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json={"prompt_id": "uuid-batch-1"}),
|
||||
httpx.Response(200, json={"prompt_id": "uuid-batch-2"}),
|
||||
]
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
|
||||
return_value=httpx.Response(200, json=queue_empty)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/history/uuid-batch-1").mock(
|
||||
return_value=httpx.Response(200, json=mock_history_1)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/history/uuid-batch-2").mock(
|
||||
return_value=httpx.Response(200, json=mock_history_2)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/view").mock(
|
||||
return_value=httpx.Response(200, content=sample_image_bytes)
|
||||
)
|
||||
|
||||
result = await generate_image(
|
||||
prompt="a landscape",
|
||||
count=2,
|
||||
seed=100,
|
||||
output_dir=str(tmp_path),
|
||||
)
|
||||
|
||||
# 4 content items: [Text1, Image1, Text2, Image2]
|
||||
assert len(result) == 4
|
||||
assert result[0].type == "text"
|
||||
assert result[1].type == "image"
|
||||
assert result[2].type == "text"
|
||||
assert result[3].type == "image"
|
||||
|
||||
# 2 files saved
|
||||
saved_files = list(tmp_path.glob("*.png"))
|
||||
assert len(saved_files) == 2
|
||||
|
||||
# Label contains batch index
|
||||
assert "1/2" in result[0].text
|
||||
assert "2/2" in result[2].text
|
||||
|
||||
|
||||
@respx.mock
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_count_2_fixed_seed_increments(
|
||||
tmp_path, sample_image_bytes, queue_empty, monkeypatch
|
||||
):
|
||||
"""count=2 with fixed seed → seeds are incremented (seed, seed+1)."""
|
||||
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path))
|
||||
|
||||
submitted_seeds = []
|
||||
|
||||
def capture_prompt(request):
|
||||
body = json.loads(request.content)
|
||||
seed_val = body["prompt"]["13"]["inputs"]["seed"]
|
||||
submitted_seeds.append(seed_val)
|
||||
idx = len(submitted_seeds)
|
||||
return httpx.Response(200, json={"prompt_id": f"seed-test-{idx}"})
|
||||
|
||||
mock_history_1 = {
|
||||
"seed-test-1": {
|
||||
"outputs": {"9": {"images": [{"filename": "img1.png", "subfolder": "", "type": "output"}]}},
|
||||
"status": {"completed": True},
|
||||
}
|
||||
}
|
||||
mock_history_2 = {
|
||||
"seed-test-2": {
|
||||
"outputs": {"9": {"images": [{"filename": "img2.png", "subfolder": "", "type": "output"}]}},
|
||||
"status": {"completed": True},
|
||||
}
|
||||
}
|
||||
|
||||
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(side_effect=capture_prompt)
|
||||
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
|
||||
return_value=httpx.Response(200, json=queue_empty)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/history/seed-test-1").mock(
|
||||
return_value=httpx.Response(200, json=mock_history_1)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/history/seed-test-2").mock(
|
||||
return_value=httpx.Response(200, json=mock_history_2)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/view").mock(
|
||||
return_value=httpx.Response(200, content=sample_image_bytes)
|
||||
)
|
||||
|
||||
await generate_image(
|
||||
prompt="a test",
|
||||
count=2,
|
||||
seed=42,
|
||||
output_dir=str(tmp_path),
|
||||
)
|
||||
|
||||
assert submitted_seeds == [42, 43], f"Expected [42, 43], got {submitted_seeds}"
|
||||
|
||||
|
||||
@respx.mock
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_count_partial_failure_continues(
|
||||
tmp_path, sample_image_bytes, queue_empty, monkeypatch
|
||||
):
|
||||
"""count=2 where first image fails → error in slot 1, second image succeeds in slot 2."""
|
||||
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path))
|
||||
|
||||
mock_history_2 = {
|
||||
"uuid-ok": {
|
||||
"outputs": {"9": {"images": [{"filename": "img2.png", "subfolder": "", "type": "output"}]}},
|
||||
"status": {"completed": True},
|
||||
}
|
||||
}
|
||||
|
||||
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
|
||||
side_effect=[
|
||||
httpx.Response(500, json={"error": "GPU OOM"}), # first fails
|
||||
httpx.Response(200, json={"prompt_id": "uuid-ok"}), # second succeeds
|
||||
]
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
|
||||
return_value=httpx.Response(200, json=queue_empty)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/history/uuid-ok").mock(
|
||||
return_value=httpx.Response(200, json=mock_history_2)
|
||||
)
|
||||
respx.get(f"{COMFYUI_BASE}/api/view").mock(
|
||||
return_value=httpx.Response(200, content=sample_image_bytes)
|
||||
)
|
||||
|
||||
result = await generate_image(
|
||||
prompt="a test",
|
||||
count=2,
|
||||
seed=10,
|
||||
output_dir=str(tmp_path),
|
||||
)
|
||||
|
||||
# First: error TextContent only (no ImageContent)
|
||||
# Second: [TextContent, ImageContent]
|
||||
assert len(result) == 3
|
||||
assert result[0].type == "text"
|
||||
assert "500" in result[0].text or "error" in result[0].text.lower()
|
||||
assert result[1].type == "text"
|
||||
assert result[2].type == "image"
|
||||
|
||||
# Only 1 file saved (the successful one)
|
||||
saved_files = list(tmp_path.glob("*.png"))
|
||||
assert len(saved_files) == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# generate_image — existing tests (kept intact)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@respx.mock
|
||||
|
||||
Reference in New Issue
Block a user