"""Run generate_image across all checkpoint/VAE/LoRA/scheduler combinations. For each pattern, this script writes a temporary TOML config, starts the ``generate-image-sd15-mcp`` server with ``streamable-http`` transport, calls ``generate_image`` through the HTTP endpoint, saves the PNG, and records the result. Usage:: uv run python script/smoke_matrix.py --output-dir data/outputs/matrix """ from __future__ import annotations import argparse import asyncio import base64 import json import subprocess import sys import time from pathlib import Path from mcp import Client from mcp.types import ImageContent DEFAULT_PROMPT = "a small red cube on a white table, studio photograph, detailed" DEFAULT_NEGATIVE = "blurry, low quality" PATTERNS: list[dict[str, object]] = [ { "name": "01_v15_base", "checkpoint": "v1-5-pruned-emaonly.safetensors", "vae": None, "loras": [], "scheduler": "dpmpp_2m_karras", "steps": 20, "cfg_scale": 7.0, "prompt": DEFAULT_PROMPT, "negative": DEFAULT_NEGATIVE, }, { "name": "02_animekawa_base", "checkpoint": "animekawa_v10.safetensors", "vae": None, "loras": [], "scheduler": "dpmpp_2m_karras", "steps": 20, "cfg_scale": 7.0, "prompt": DEFAULT_PROMPT, "negative": DEFAULT_NEGATIVE, }, { "name": "03_amore_base", "checkpoint": "animeModel_amore.safetensors", "vae": None, "loras": [], "scheduler": "dpmpp_2m_karras", "steps": 20, "cfg_scale": 7.0, "prompt": DEFAULT_PROMPT, "negative": DEFAULT_NEGATIVE, }, { "name": "04_animepastel_base", "checkpoint": "animePastelDream_softBakedVae_full_fp16.safetensors", "vae": None, "loras": [], "scheduler": "dpmpp_2m_karras", "steps": 20, "cfg_scale": 7.0, "prompt": DEFAULT_PROMPT, "negative": DEFAULT_NEGATIVE, }, { "name": "05_v15_external_vae", "checkpoint": "v1-5-pruned-emaonly.safetensors", "vae": "vaeKlF8Anime2_klF8Anime2VAE.safetensors", "loras": [], "scheduler": "dpmpp_2m_karras", "steps": 20, "cfg_scale": 7.0, "prompt": DEFAULT_PROMPT, "negative": DEFAULT_NEGATIVE, }, { "name": "06_v15_lora_pixel_sprites", "checkpoint": "v1-5-pruned-emaonly.safetensors", "vae": None, "loras": [{"file": "pixel sprites.safetensors", "weight": 0.8}], "scheduler": "dpmpp_2m_karras", "steps": 20, "cfg_scale": 7.0, "prompt": "pixel art, a small red cube on a white table, detailed", "negative": DEFAULT_NEGATIVE, }, { "name": "07_v15_lora_pixelartV3", "checkpoint": "v1-5-pruned-emaonly.safetensors", "vae": None, "loras": [{"file": "pixelartV3.safetensors", "weight": 0.8}], "scheduler": "dpmpp_2m_karras", "steps": 20, "cfg_scale": 7.0, "prompt": "pixelart, a small red cube on a white table, detailed", "negative": DEFAULT_NEGATIVE, }, { "name": "08_v15_lora_pixel_f2", "checkpoint": "v1-5-pruned-emaonly.safetensors", "vae": None, "loras": [{"file": "pixel_f2.safetensors", "weight": 0.8}], "scheduler": "dpmpp_2m_karras", "steps": 20, "cfg_scale": 7.0, "prompt": "pixel, a small red cube on a white table, detailed", "negative": DEFAULT_NEGATIVE, }, { "name": "09_v15_lcm_lora", "checkpoint": "v1-5-pruned-emaonly.safetensors", "vae": None, "loras": [{"file": "LCM_LoRA_Weights_SD15.safetensors", "weight": 1.0}], "scheduler": "lcm", "steps": 4, "cfg_scale": 1.0, "prompt": DEFAULT_PROMPT, "negative": DEFAULT_NEGATIVE, }, { "name": "10_animekawa_vae_lora", "checkpoint": "animekawa_v10.safetensors", "vae": "vaeKlF8Anime2_klF8Anime2VAE.safetensors", "loras": [{"file": "pixelartV3.safetensors", "weight": 0.7}], "scheduler": "euler_a", "steps": 20, "cfg_scale": 7.0, "prompt": "pixelart, a small red cube on a white table, detailed", "negative": DEFAULT_NEGATIVE, }, ] def make_toml(pattern: dict[str, object], port: int) -> str: lines: list[str] = [] lines.append("[model]") lines.append('root = "."') lines.append(f'checkpoint = "{pattern["checkpoint"]}"') if pattern["vae"]: lines.append(f'vae = "{pattern["vae"]}"') for lora in pattern["loras"]: # type: ignore[union-attr] lines.append("[[model.lora]]") lines.append(f'file = "{lora["file"]}"') # type: ignore[index] lines.append(f"weight = {lora['weight']}") # type: ignore[index] lines.append("[generation]") lines.append(f"steps = {pattern['steps']}") lines.append(f"cfg_scale = {pattern['cfg_scale']}") lines.append(f'scheduler = "{pattern["scheduler"]}"') lines.append('default_size = "square_512"') lines.append("[output]") lines.append("enabled = false") lines.append("[server]") lines.append('transport = "streamable-http"') lines.append('host = "127.0.0.1"') lines.append(f"port = {port}") return "\n".join(lines) + "\n" async def run_one( pattern: dict[str, object], config_path: Path, output_path: Path, port: int, ) -> dict[str, object]: toml = make_toml(pattern, port) config_path.write_text(toml, encoding="utf-8") proc = subprocess.Popen( [sys.executable, "-m", "mcp_sd15", "--config", str(config_path)], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, ) result: dict[str, object] = {"name": pattern["name"], "status": "unknown"} started = time.monotonic() try: # Wait for server to be ready (up to 120s for model loading) url = f"http://127.0.0.1:{port}/mcp" for _wait in range(120): if proc.poll() is not None: # Process died; drain stdout dead_output = "" if proc.stdout: dead_output = proc.stdout.read() result["status"] = "server_died" result["output"] = dead_output[-3000:] return result try: import httpx resp = httpx.get(url, timeout=2.0) if resp.status_code == 400: break except Exception: pass await asyncio.sleep(1) else: result["status"] = "server_timeout" return result async with Client(url, raise_exceptions=True) as client: tools = (await client.list_tools()).tools tool_names = {tool.name for tool in tools} if tool_names != {"generate_image", "generate_guided_image", "inpaint_image"}: result["status"] = "wrong_tools" result["tools"] = list(tool_names) return result call_result = await client.call_tool( "generate_image", { "positive_prompt": pattern["prompt"], # type: ignore[index] "negative_prompt": pattern["negative"], # type: ignore[index] "size": "square_512", }, ) if call_result.is_error or len(call_result.content) != 1: result["status"] = "tool_error" result["detail"] = str(call_result) return result content = call_result.content[0] if not isinstance(content, ImageContent) or content.mime_type != "image/png": result["status"] = "wrong_content" result["detail"] = str(content) return result output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_bytes(base64.b64decode(content.data, validate=True)) elapsed = round(time.monotonic() - started, 1) result["status"] = "ok" result["elapsed"] = elapsed result["output_file"] = str(output_path) return result except Exception as exc: result["status"] = "exception" result["detail"] = str(exc) return result finally: proc.terminate() try: proc.wait(timeout=10) except subprocess.TimeoutExpired: proc.kill() async def main(args: argparse.Namespace) -> None: config_path = args.data_dir / "config.matrix.toml" results: list[dict[str, object]] = [] for i, pattern in enumerate(PATTERNS): name = pattern["name"] print(f"[{i + 1}/{len(PATTERNS)}] {name} ...", flush=True) output_path = args.output_dir / f"{name}.png" result = await run_one(pattern, config_path, output_path, args.port + i) status = result["status"] if status == "ok": print( f" -> OK elapsed={result.get('elapsed')}s file={result.get('output_file')}", flush=True, ) else: print(f" -> FAIL status={status} detail={result.get('detail', '')}", flush=True) results.append(result) config_path.unlink(missing_ok=True) summary_path = args.output_dir / "matrix_results.json" summary_path.parent.mkdir(parents=True, exist_ok=True) summary_path.write_text(json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8") ok = sum(1 for r in results if r["status"] == "ok") fail = len(results) - ok print(f"\n=== Summary: {ok} ok / {fail} fail / {len(results)} total ===", flush=True) if fail: print("Failures:", flush=True) for r in results: if r["status"] != "ok": print(f" {r['name']}: {r['status']} {r.get('detail', '')}", flush=True) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--output-dir", type=Path, default=Path("data/outputs/matrix")) parser.add_argument("--data-dir", type=Path, default=Path("data")) parser.add_argument("--port", type=int, default=8130) return parser.parse_args() if __name__ == "__main__": asyncio.run(main(parse_args()))