From 79f1e6d65f2c3e85fc0a61b4629b704b4e5b3934 Mon Sep 17 00:00:00 2001 From: Patrick Plate Date: Mon, 6 Apr 2026 07:45:37 +0200 Subject: [PATCH] feat(mcp-image-gen): add name and count params to generate_image - Add name (str) param: filename prefix saved as {name}_{timestamp}_{seed}.png - Add count (int, 1-10) param: generate N images in one call - Extract _sanitize_name() helper: strips special chars, collapses underscores, caps at 64 chars - Extract _build_filename() helper: pure function for testable filename construction - Extract _generate_single() coroutine: clean loop body for batch generation - Fixed seed batches increment seed per image (seed+i-1) for deterministic variation - random seed (-1) batches give independent random seeds per image - Partial batch failures continue (error TextContent in slot, remaining images proceed) - Returns flat interleaved [Text1, Image1, Text2, Image2, ...] list - 34/34 tests passing (was 19, added 15 new tests) --- .../EXPAND_GENERATE_IMAGE_Assessment.md | 319 +++++++++++++++++ mcp/mcp-image-gen/src/server.py | 200 ++++++++--- mcp/mcp-image-gen/tests/test_server.py | 320 +++++++++++++++++- 3 files changed, 794 insertions(+), 45 deletions(-) create mode 100644 mcp/mcp-image-gen/EXPAND_GENERATE_IMAGE_Assessment.md diff --git a/mcp/mcp-image-gen/EXPAND_GENERATE_IMAGE_Assessment.md b/mcp/mcp-image-gen/EXPAND_GENERATE_IMAGE_Assessment.md new file mode 100644 index 0000000..2dcd8bb --- /dev/null +++ b/mcp/mcp-image-gen/EXPAND_GENERATE_IMAGE_Assessment.md @@ -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 | diff --git a/mcp/mcp-image-gen/src/server.py b/mcp/mcp-image-gen/src/server.py index 9920525..b55ed65 100644 --- a/mcp/mcp-image-gen/src/server.py +++ b/mcp/mcp-image-gen/src/server.py @@ -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. diff --git a/mcp/mcp-image-gen/tests/test_server.py b/mcp/mcp-image-gen/tests/test_server.py index 2fd1785..0707f11 100644 --- a/mcp/mcp-image-gen/tests/test_server.py +++ b/mcp/mcp-image-gen/tests/test_server.py @@ -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