diff options
Diffstat (limited to 'script')
| -rw-r--r-- | script/benchmark_generate_image.py | 198 |
1 files changed, 198 insertions, 0 deletions
diff --git a/script/benchmark_generate_image.py b/script/benchmark_generate_image.py new file mode 100644 index 0000000..8ba74a5 --- /dev/null +++ b/script/benchmark_generate_image.py @@ -0,0 +1,198 @@ +"""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) + 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() |