summaryrefslogtreecommitdiff
path: root/script/smoke_matrix.py
diff options
context:
space:
mode:
Diffstat (limited to 'script/smoke_matrix.py')
-rw-r--r--script/smoke_matrix.py304
1 files changed, 304 insertions, 0 deletions
diff --git a/script/smoke_matrix.py b/script/smoke_matrix.py
new file mode 100644
index 0000000..7fc3ec2
--- /dev/null
+++ b/script/smoke_matrix.py
@@ -0,0 +1,304 @@
+"""Run generate_image across all checkpoint/VAE/LoRA/scheduler combinations.
+
+For each pattern, this script writes a temporary TOML config, starts the
+``mcp-sd15`` 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()))
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-19 14:27:21 +0000