"""Benchmark repeatable generate_image calls through the in-process MCP server.""" from __future__ import annotations import argparse import asyncio import base64 import io import json import time from dataclasses import asdict, dataclass from pathlib import Path from unittest.mock import patch import numpy as np import torch from mcp import Client from mcp.types import ImageContent from PIL import Image from mcp_sd15.config import load_config from mcp_sd15.model_manager import ModelManager from mcp_sd15.server import create_server from mcp_sd15.service import ImageGenerationService @dataclass(frozen=True) class Scene: name: str seed: int size: str prompt: str SCENES = ( Scene( name="full_body_white", seed=2026091201, size="portrait_512", prompt=( "1girl, solo, cute adult girl, full body, front view, pink bob hair, blue eyes, " "blue dress, hands visible, feet visible, centered, solid white background, no props" ), ), Scene( name="dynamic_pose_white", seed=2026091202, size="square_512", prompt=( "1girl, solo, cute adult girl, three-quarter body, pink bob hair, blue eyes, " "blue dress, waving hand, five fingers, other hand on hip, centered, " "solid white background, no props" ), ), ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, required=True) parser.add_argument("--label", required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--result", type=Path, required=True) return parser.parse_args() def mib(value: int) -> float: return round(value / 1024**2, 1) def image_metrics(png_data: bytes) -> dict[str, float]: with Image.open(io.BytesIO(png_data)) as image: rgb = np.asarray(image.convert("RGB"), dtype=np.float32) height, width, _ = rgb.shape border = max(8, min(height, width) // 16) border_pixels = np.concatenate( ( rgb[:border].reshape(-1, 3), rgb[-border:].reshape(-1, 3), rgb[border:-border, :border].reshape(-1, 3), rgb[border:-border, -border:].reshape(-1, 3), ) ) white = np.all(rgb >= 245, axis=2) border_white = np.all(border_pixels >= 245, axis=1) luminance = rgb.mean(axis=2) edge_strength = ( np.abs(np.diff(luminance, axis=0)).mean() + np.abs(np.diff(luminance, axis=1)).mean() ) / 2 return { "white_pixel_ratio": round(float(white.mean()), 4), "border_white_ratio": round(float(border_white.mean()), 4), "border_rgb_std": round(float(border_pixels.std()), 2), "edge_strength": round(float(edge_strength), 2), } async def run(args: argparse.Namespace) -> None: if not torch.cuda.is_available(): raise RuntimeError("CUDA is required for this benchmark") args.output_dir.mkdir(parents=True, exist_ok=True) args.result.parent.mkdir(parents=True, exist_ok=True) torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() free_before, total = torch.cuda.mem_get_info() config = load_config(args.config) load_started = time.perf_counter() backend = ModelManager(config) backend.load() torch.cuda.synchronize() load_seconds = time.perf_counter() - load_started model_allocated = torch.cuda.memory_allocated() model_reserved = torch.cuda.memory_reserved() scene_reports: list[dict[str, object]] = [] report: dict[str, object] = { "label": args.label, "config": str(args.config.resolve()), "checkpoint": config.model.checkpoint, "vae": config.model.vae, "loras": [asdict(item) for item in config.model.loras], "scheduler": config.generation.scheduler, "steps": config.generation.steps, "cfg_scale": config.generation.cfg_scale, "gpu": torch.cuda.get_device_name(), "gpu_total_mib": mib(total), "gpu_free_before_mib": mib(free_before), "load_seconds": round(load_seconds, 3), "model_allocated_mib": mib(model_allocated), "model_reserved_mib": mib(model_reserved), "scenes": scene_reports, } service = ImageGenerationService(config, backend) server = create_server(service) async with Client(server, raise_exceptions=True) as client: for scene in SCENES: positive = f"{config.prompt.positive}, {scene.prompt}" negative = config.prompt.negative tokenizer = backend.base.tokenizer positive_tokens = len(tokenizer(positive, add_special_tokens=True).input_ids) negative_tokens = len(tokenizer(negative, add_special_tokens=True).input_ids) token_limit = tokenizer.model_max_length if positive_tokens > token_limit or negative_tokens > token_limit: raise RuntimeError( f"prompt exceeds tokenizer limit {token_limit}: " f"positive={positive_tokens}, negative={negative_tokens}" ) torch.cuda.reset_peak_memory_stats() started = time.perf_counter() with patch("mcp_sd15.service.secrets.randbelow", return_value=scene.seed): result = await client.call_tool( "generate_image", { "positive_prompt": scene.prompt, "negative_prompt": None, "size": scene.size, }, ) torch.cuda.synchronize() elapsed = time.perf_counter() - started if result.is_error or len(result.content) != 1: raise RuntimeError(f"generate_image failed: {result}") content = result.content[0] if not isinstance(content, ImageContent) or content.mime_type != "image/png": raise RuntimeError(f"unexpected MCP content: {content}") png_data = base64.b64decode(content.data, validate=True) output = args.output_dir / f"{args.label}__{scene.name}.png" output.write_bytes(png_data) scene_report = { "name": scene.name, "seed": scene.seed, "size": scene.size, "positive_tokens": positive_tokens, "negative_tokens": negative_tokens, "elapsed_seconds": round(elapsed, 3), "peak_allocated_mib": mib(torch.cuda.max_memory_allocated()), "peak_reserved_mib": mib(torch.cuda.max_memory_reserved()), "output": str(output.resolve()), **image_metrics(png_data), } scene_reports.append(scene_report) print(json.dumps(scene_report, ensure_ascii=False)) args.result.write_text( json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" ) print(f"result={args.result.resolve()}") def main() -> None: asyncio.run(run(parse_args())) if __name__ == "__main__": main()