1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
|
"""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()
|
