path: root/script/smoke_matrix.py

"""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()))