summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--README.md21
-rw-r--r--docs/lazy-loading-verification.md75
-rw-r--r--script/benchmark_generate_image.py1
-rw-r--r--script/smoke_stdio.py153
-rw-r--r--src/mcp_sd15/cli.py3
-rw-r--r--src/mcp_sd15/model_manager.py36
-rw-r--r--tests/test_lazy_loading.py245
7 files changed, 530 insertions, 4 deletions
diff --git a/README.md b/README.md
index 50f98ec..bf5e6ac 100644
--- a/README.md
+++ b/README.md
@@ -40,7 +40,7 @@ positiveは`base, user`の順でカンマ区切り、negativeは`base\nuser`の
モデルは `<root>/checkpoints`、`vae`、`lora`、`controlnet` からだけ読み込みます。
VAE componentは生成に常に必要です。`model.vae`を省略した場合はcheckpoint内蔵VAEを使い、
-指定した場合だけ外部VAEへ置き換えます。設定されたLoRAは起動時にすべてロードされ、3つの
+指定した場合だけ外部VAEへ置き換えます。設定されたLoRAは初回生成時にすべてロードされ、3つの
生成ツールに共通して適用されます。MCP呼び出しごとのLoRA切替には対応しません。
Depth/Openpose の annotator 重みは `controlnet_aux` のローカルキャッシュに事前配置して
ください。推論時のダウンロードは行いません。
@@ -53,6 +53,25 @@ GPU を使わないテストは `uv run pytest`、lint は `uv run ruff check .`
`uv run pyright` で実行できます。実モデルの smoke test は環境依存のため自動テストには
含めていません。
+起動・MCP初期化・ツール一覧取得では画像モデルをロードせず、CUDAを初期化しません。
+設定とモデルファイルの存在は起動時に確認し、有効な生成要求の実行時にSD本体・VAE・LoRAを
+ワーカースレッドでロードします。モデル互換性やCUDAのエラーは初回生成時に返されます。
+同一プロセスではロード済みモデルを再利用し、stdioプロセスの終了で解放します。
+一覧取得と生成を別プロセスで行うクライアントでは、生成プロセスごとのロードが必要です。
+
+実際のCLI/stdio経路の一覧取得とVRAMを確認するには次を実行します（常駐化は不要です）。
+
+```powershell
+uv run python script\smoke_stdio.py --config data\config.toml `
+ --result data\verification-lazy-load\list.json
+```
+
+`--images 2`を追加すると、固定seedで同じ接続から2枚生成し、PNG、画素ハッシュ、時間、
+PNGメタデータ、stderr、生成前後とプロセス終了後の`nvidia-smi`結果を記録します。
+seed固定は確認スクリプト内だけに適用され、製品のMCPスキーマやseed処理は変えません。
+要求の待機上限は`--timeout`で指定します（既定300秒）。Windowsでプロセス別VRAMが
+`N/A`の場合は取得不能として扱い、GPU全体の使用量とプロセスの終了も確認してください。
+
実GPUとMCP経路を使う`generate_image` smoke testは次のように実行できます。
```powershell
diff --git a/docs/lazy-loading-verification.md b/docs/lazy-loading-verification.md
new file mode 100644
index 0000000..06b82bb
--- /dev/null
+++ b/docs/lazy-loading-verification.md
@@ -0,0 +1,75 @@
+# 画像モデル遅延ロードの検証結果
+
+実施日: 2026-10-14。常駐化やクライアントの接続方式は変更せず、モデル初期化のタイミングだけを変更した。
+
+## 回帰テスト
+
+| 検査 | 修正前 | 修正後 |
+| --- | --- | --- |
+| `python -m pytest -q` | 69 passed | 87 passed |
+| `ruff check .` | 成功 | 成功 |
+| `pyright` | 0 errors, 0 warnings | 0 errors, 0 warnings |
+
+既存69件のテストと期待値は変更していない。追加18件で、一覧取得時の推論ライブラリの
+import禁止、CLI/stdio経路、全生成方式の初回ロード、ワーカースレッドでの実行、
+同時要求・明示的な同時ロード、無効な入力、ロード失敗後の再試行、ロードログを確認した。
+推論依存ファイルがサンドボックス内で読み取り拒否となったため、実行時は通常の権限を使用した。
+
+## 実GPU・stdio経路
+
+現在の `data/config.toml` を使用した。checkpoint、外部VAE、2件のLoRA、LCM、6 stepsは変更していない。
+seedは確認用の子プロセス内だけで `2026091201` に固定した。
+
+| 計測 | 修正前 | 修正後 |
+| --- | --- | --- |
+| 一覧取得のみ: 起動～初期化～一覧取得 | 8.233秒 | 0.671秒 |
+| 生成セッション: 起動～初期化～一覧取得 | 7.898秒 | 0.650秒 |
+| 生成セッション: 初回生成要求～PNG応答 | 17.656秒 | 25.089秒 |
+| 生成セッション: 2枚目の生成要求～PNG応答 | 1.211秒 | 1.229秒 |
+| 一覧取得のみ: GPU全体の使用量（開始→取得後） | 808→3233 MiB | 682→682 MiB |
+| 生成セッション: GPU全体の使用量（開始→終了後） | 806→674 MiB | 682→692 MiB |
+
+各値は単一セッションの観測値で、統計的な速度保証ではない。初回生成にはCUDA等のwarm-upも含まれる。
+修正後のベースモデルのロード時間は7.812秒。一覧取得時のロードは0回、同じ接続で2枚生成しても
+ロード開始・完了はそれぞれ1回で、生成開始後にのみ記録された。
+
+最終版の確認スクリプトでも2枚生成して再確認した。一覧取得は0.655秒、初回生成は14.089秒、
+2枚目は1.221秒だった。画素・保存PNGは一致し、ロードは1回、生成プロセスの終了も確認した。
+この再実行では一覧取得前後のGPU全体の使用量は683→683 MiB、終了後は693 MiBだった。
+
+Windowsのプロセス別VRAMは `N/A` だったため、個別使用量は測定不能。
+GPU全体の使用量には他プロセスの変動が含まれる。修正前後の一覧取得・生成の計4プロセスは
+`tasklist` で終了を確認し、生成プロセスは終了後のGPUプロセス一覧にも残らなかった。
+
+## 画像生成の回帰確認
+
+修正前後で同じ設定・プロンプト・seedによる2枚ずつの画像を比較した。
+
+- すべて512×512のPNGで、修正前後の画素は完全一致した。
+- 画素SHA-256は `5325bca0f3bf757e40ef28dfd4e6e16d0d316a95285f770cef7b55b9d3d92137`。
+- PNGメタデータのキーとseed・size・mode・schedulerの値は一致した。
+- timestamp・elapsed_secondsは実行時刻・ロード時間を反映して変わることを確認した。
+- 修正前後とも返却されたPNGとアプリが保存したPNGがバイト単位で一致した。
+
+ControlNetの実GPU生成は現在の設定で無効のため未実施。
+誘導生成・Inpaint・前処理のロード順はスタブを用いた自動テストで確認した。
+
+## 通常質問の経路と確認範囲
+
+親アプリの `_chat_event_stream` に通常の質問を渡し、実際の設定済みuvコマンドを使った
+MCP起動・一覧取得・3ツールの登録を確認した。ロードログは発生せず、GPU全体の使用量は
+693→682 MiBだった。他のMCPを起動しないよう対象をこのサーバーに限定し、allowlistは
+確認プロセス内だけに設定した。会話の保存は行っていない。
+
+LLM応答はスタブ化した。ブラウザ操作と実LLMを含むチャット全体の確認は未実施。
+
+## 再実行と成果物
+
+```powershell
+.venv/Scripts/python.exe script/smoke_stdio.py --config data/config.toml --result data/verification-lazy-load/list.json
+.venv/Scripts/python.exe script/smoke_stdio.py --config data/config.toml --result data/verification-lazy-load/generate.json --images 2
+```
+
+生の計測JSON・stderr・PNG・比較結果・チャット確認結果は
+`data/verification-lazy-load/` に保存した（既存のgitignore対象）。
+初回生成のロード時間と、別プロセスで生成する際の再ロードは仕様として維持している。
diff --git a/script/benchmark_generate_image.py b/script/benchmark_generate_image.py
index 8ba74a5..2752d6d 100644
--- a/script/benchmark_generate_image.py
+++ b/script/benchmark_generate_image.py
@@ -108,6 +108,7 @@ async def run(args: argparse.Namespace) -> None:
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()
diff --git a/script/smoke_stdio.py b/script/smoke_stdio.py
new file mode 100644
index 0000000..9f0dfab
--- /dev/null
+++ b/script/smoke_stdio.py
@@ -0,0 +1,153 @@
+"""Measure CLI/stdio discovery and optional real generation without a resident server."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import hashlib
+import io
+import json
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+from mcp import ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from PIL import Image
+
+
+def gpu_snapshot() -> dict[str, object]:
+ """Observe externally; do not initialize CUDA in the client or server."""
+ result: dict[str, object] = {}
+ for name, query in (
+ ("gpu", "--query-gpu=index,memory.used,memory.total"),
+ ("processes", "--query-compute-apps=pid,used_gpu_memory"),
+ ):
+ try:
+ completed = subprocess.run(
+ ["nvidia-smi", query, "--format=csv,noheader,nounits"],
+ capture_output=True,
+ text=True,
+ timeout=10,
+ check=True,
+ )
+ result[name] = completed.stdout.strip().splitlines()
+ except (OSError, subprocess.SubprocessError) as exc:
+ result[name] = {"unavailable": str(exc)}
+ return result
+
+
+async def run(args: argparse.Namespace) -> None:
+ args.result.parent.mkdir(parents=True, exist_ok=True)
+ log_path = args.result.with_suffix(".stderr.log")
+ # Fixed seeds belong to this verification process, never the production CLI/schema.
+ bootstrap = (
+ "import os,sys; print('smoke_pid='+str(os.getpid()),file=sys.stderr,flush=True); "
+ "from mcp_sd15.service import secrets; "
+ f"secrets.randbelow=lambda _: {args.seed}; "
+ "from mcp_sd15.cli import main; raise SystemExit(main())"
+ )
+ params = StdioServerParameters(
+ command=sys.executable,
+ args=["-c", bootstrap, "--config", str(args.config.resolve())],
+ cwd=str(Path(__file__).resolve().parents[1]),
+ )
+ report: dict[str, object] = {
+ "config": str(args.config.resolve()),
+ "seed": args.seed,
+ "gpu_before": await asyncio.to_thread(gpu_snapshot),
+ "stderr": str(log_path.resolve()),
+ }
+ try:
+ with log_path.open("w", encoding="utf-8") as errlog:
+ started = time.perf_counter()
+ async with stdio_client(params, errlog=errlog) as (read, write):
+ report["spawn_seconds"] = round(time.perf_counter() - started, 3)
+ async with ClientSession(
+ read, write, read_timeout_seconds=args.timeout
+ ) as client:
+ initialized = time.perf_counter()
+ await client.initialize()
+ report["initialize_seconds"] = round(
+ time.perf_counter() - initialized, 3
+ )
+ listed = time.perf_counter()
+ tools = (await client.list_tools()).tools
+ report["list_seconds"] = round(time.perf_counter() - listed, 3)
+ report["discovery_seconds"] = round(time.perf_counter() - started, 3)
+ names = {tool.name for tool in tools}
+ if names != {"generate_image", "generate_guided_image", "inpaint_image"}:
+ raise RuntimeError(f"unexpected tools: {names}")
+ report["tools"] = sorted(names)
+ report["gpu_after_list"] = await asyncio.to_thread(gpu_snapshot)
+ report["stderr_after_list"] = log_path.read_text(encoding="utf-8")
+ images: list[dict[str, object]] = []
+ report["images"] = images
+ for index in range(args.images):
+ generated = time.perf_counter()
+ result = await client.call_tool(
+ "generate_image",
+ {
+ "positive_prompt": args.prompt,
+ "negative_prompt": "blurry, low quality",
+ "size": "square_512",
+ },
+ )
+ elapsed = time.perf_counter() - generated
+ if getattr(result, "is_error", getattr(result, "isError", False)):
+ raise RuntimeError(f"generation failed: {result}")
+ if len(result.content) != 1 or result.content[0].type != "image":
+ raise RuntimeError(f"unexpected content: {result}")
+ content = result.content[0]
+ if getattr(content, "mime_type", getattr(content, "mimeType", None)) != (
+ "image/png"
+ ):
+ raise RuntimeError("generation did not return image/png")
+ png = base64.b64decode(content.data, validate=True)
+ destination = args.result.with_name(f"{args.result.stem}-{index + 1}.png")
+ destination.write_bytes(png)
+ with Image.open(io.BytesIO(png)) as image:
+ image.load()
+ if image.format != "PNG" or image.size != (512, 512):
+ raise RuntimeError("generation returned unexpected format or size")
+ images.append(
+ {
+ "elapsed_seconds": round(elapsed, 3),
+ "format": image.format,
+ "dimensions": image.size,
+ "metadata": image.info,
+ "pixel_sha256": hashlib.sha256(image.tobytes()).hexdigest(),
+ "output": str(destination.resolve()),
+ "gpu": await asyncio.to_thread(gpu_snapshot),
+ }
+ )
+ report["stderr_after_calls"] = log_path.read_text(encoding="utf-8")
+ report["status"] = "passed"
+ except BaseException as exc:
+ report["status"] = "failed"
+ report["error"] = repr(exc)
+ raise
+ finally:
+ report["gpu_after_exit"] = await asyncio.to_thread(gpu_snapshot)
+ report["stderr_final"] = log_path.read_text(encoding="utf-8")
+ 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:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--config", type=Path, required=True)
+ parser.add_argument("--result", type=Path, required=True)
+ parser.add_argument("--images", type=int, choices=(0, 1, 2), default=0)
+ parser.add_argument("--seed", type=int, default=2026091201)
+ parser.add_argument("--timeout", type=float, default=300)
+ parser.add_argument("--prompt", default="a small red cube on a white table, studio photograph")
+ asyncio.run(run(parser.parse_args()))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/mcp_sd15/cli.py b/src/mcp_sd15/cli.py
index abfde01..1248581 100644
--- a/src/mcp_sd15/cli.py
+++ b/src/mcp_sd15/cli.py
@@ -31,13 +31,14 @@ def main(argv: Sequence[str] | None = None) -> int:
try:
config = load_config(args.config)
logging.getLogger(__name__).info(
- "loading checkpoint=%s loras=%d",
+ "configuration loaded checkpoint=%s loras=%d; models load on generation",
config.model.checkpoint_path(),
len(config.model.loras),
)
backend = ModelManager(config)
service = ImageGenerationService(config, backend)
server = create_server(service)
+ logging.getLogger(__name__).info("starting MCP transport=%s", config.server.transport)
if config.server.transport == "streamable-http":
server.run(
transport="streamable-http",
diff --git a/src/mcp_sd15/model_manager.py b/src/mcp_sd15/model_manager.py
index e87a813..2c314d1 100644
--- a/src/mcp_sd15/model_manager.py
+++ b/src/mcp_sd15/model_manager.py
@@ -3,6 +3,9 @@
from __future__ import annotations
import gc
+import logging
+import threading
+import time
from typing import Any
import numpy as np
@@ -15,21 +18,48 @@ from .preprocessors.depth import DepthPreprocessor
from .preprocessors.openpose import OpenPosePreprocessor
from .schema import GuidedImageRequest, InpaintImageRequest, TextToImageRequest
+logger = logging.getLogger(__name__)
+
class ModelManager:
- """Own all model state; construction validates model compatibility at startup."""
+ """Own model state, loading inference models only when generation needs them."""
def __init__(self, config: AppConfig) -> None:
self.config = config
self.factory = PipelineFactory(config)
- self.base = self.factory.create_base()
+ self.base: Any = None
+ self._load_lock = threading.Lock()
self._control_mode: str | None = None
self._controlnet: Any | None = None
self._control_pipeline: Any | None = None
self._depth = DepthPreprocessor()
self._openpose = OpenPosePreprocessor()
+ def load(self) -> None:
+ """Load once on success; failed initialization remains retryable.
+
+ Generation calls this from the service's serialized worker thread.
+ The lock also protects explicit loads by measurement scripts.
+ """
+ with self._load_lock:
+ if self.base is not None:
+ return
+ started = time.monotonic()
+ logger.info(
+ "model load started checkpoint=%s loras=%d",
+ self.config.model.checkpoint_path(),
+ len(self.config.model.loras),
+ )
+ try:
+ base = self.factory.create_base()
+ except Exception:
+ logger.exception("model load failed")
+ raise
+ self.base = base
+ logger.info("model load completed elapsed_seconds=%.3f", time.monotonic() - started)
+
def text_to_image(self, request: TextToImageRequest, seed: int) -> Image.Image:
+ self.load()
width, height = request.size.dimensions
prompt_embeds, negative_prompt_embeds = self._embed_prompts(
request.positive_prompt, request.negative_prompt
@@ -49,6 +79,7 @@ class ModelManager:
def guided_image(
self, request: GuidedImageRequest, image: Image.Image, seed: int
) -> Image.Image:
+ self.load()
width, height = request.size.dimensions
image = image.resize((width, height), Image.Resampling.LANCZOS)
if request.control_mode == "depth":
@@ -79,6 +110,7 @@ class ModelManager:
mask: Image.Image,
seed: int,
) -> Image.Image:
+ self.load()
width, height = request.size.dimensions
image = image.resize((width, height), Image.Resampling.LANCZOS)
mask = mask.resize((width, height), Image.Resampling.NEAREST)
diff --git a/tests/test_lazy_loading.py b/tests/test_lazy_loading.py
new file mode 100644
index 0000000..bd4b341
--- /dev/null
+++ b/tests/test_lazy_loading.py
@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import sys
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import replace
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+from mcp import Client, ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from PIL import Image
+
+from mcp_sd15.errors import ConfigError
+from mcp_sd15.model_manager import ModelManager
+from mcp_sd15.pipeline_factory import PipelineFactory
+from mcp_sd15.schema import TextToImageRequest
+from mcp_sd15.server import create_server
+from mcp_sd15.service import ImageGenerationService
+
+
+@pytest.fixture
+def lazy_backend(app_config, monkeypatch):
+ loads: list[int] = []
+ calls: list[dict[str, object]] = []
+
+ def pipeline(**kwargs):
+ calls.append(kwargs)
+ return SimpleNamespace(images=[Image.new("RGB", (4, 4), "red")])
+
+ def create_base(self):
+ loads.append(threading.get_ident())
+ return pipeline
+
+ monkeypatch.setattr(PipelineFactory, "create_base", create_base)
+ backend = ModelManager(app_config)
+ monkeypatch.setattr(backend, "_embed_prompts", lambda *args: ("positive", "negative"))
+ monkeypatch.setattr(backend, "_generator", lambda seed: seed)
+
+ def preprocess(image, **kwargs):
+ assert backend.base is pipeline # Base must load before any GPU/preprocessor work.
+ return image
+
+ monkeypatch.setattr(backend, "_depth", preprocess)
+ monkeypatch.setattr(backend, "_openpose", preprocess)
+ monkeypatch.setattr(backend, "_inpaint_condition", lambda image, mask: preprocess(image))
+ monkeypatch.setattr(backend, "_pipeline_for", lambda *args, **kwargs: pipeline)
+ return backend, loads, calls
+
+
+async def test_discovery_does_not_load_model(app_config, lazy_backend) -> None:
+ backend, loads, _ = lazy_backend
+ assert backend.base is None
+ server = create_server(ImageGenerationService(app_config, backend))
+ assert loads == []
+ async with Client(server, mode="legacy") as client:
+ tools = (await client.list_tools()).tools
+ assert len(tools) == 3
+ assert loads == []
+ assert backend.base is None
+
+
+@pytest.mark.parametrize(
+ ("tool", "arguments"),
+ [
+ ("generate_image", {}),
+ ("generate_guided_image", {"control_mode": "depth"}),
+ ("generate_guided_image", {"control_mode": "openpose"}),
+ ("generate_guided_image", {"control_mode": "openpose_full"}),
+ ("inpaint_image", {}),
+ ],
+)
+async def test_first_generation_loads_once_in_worker(
+ app_config, lazy_backend, image_base64, tool, arguments, caplog
+) -> None:
+ caplog.set_level(logging.INFO, logger="mcp_sd15.model_manager")
+ backend, loads, calls = lazy_backend
+ arguments = {"positive_prompt": "cat", **arguments}
+ if tool != "generate_image":
+ arguments["input_image"] = image_base64
+ if tool == "inpaint_image":
+ arguments["mask_image"] = image_base64
+ server = create_server(ImageGenerationService(app_config, backend))
+ async with Client(server) as client:
+ for _ in range(2):
+ result = await client.call_tool(tool, arguments)
+ assert result.is_error is False
+ assert result.content[0].type == "image"
+ assert len(loads) == 1
+ assert loads[0] != threading.get_ident()
+ assert len(calls) == 2
+ assert calls[0]["num_inference_steps"] == app_config.generation.steps
+ assert calls[0]["guidance_scale"] == app_config.generation.cfg_scale
+ assert caplog.text.count("model load started") == 1
+ assert caplog.text.count("model load completed") == 1
+
+
+async def test_simultaneous_requests_serialize_load_and_generation(
+ app_config, lazy_backend
+) -> None:
+ backend, loads, calls = lazy_backend
+ service = ImageGenerationService(app_config, backend)
+ await asyncio.gather(
+ *(service.generate_image(TextToImageRequest.create("cat")) for _ in range(4))
+ )
+ assert len(loads) == 1
+ assert len(calls) == 4
+
+
+def test_load_is_thread_safe_and_idempotent(app_config, monkeypatch) -> None:
+ loads: list[int] = []
+ pipeline = object()
+
+ def create_base(self):
+ loads.append(threading.get_ident())
+ time.sleep(0.03)
+ return pipeline
+
+ monkeypatch.setattr(PipelineFactory, "create_base", create_base)
+ backend = ModelManager(app_config)
+ with ThreadPoolExecutor(max_workers=8) as executor:
+ list(executor.map(lambda _: backend.load(), range(8)))
+ backend.load()
+ assert len(loads) == 1
+ assert backend.base is pipeline
+
+
+@pytest.mark.parametrize(
+ ("tool", "arguments"),
+ [
+ ("generate_image", {"positive_prompt": ""}),
+ ("generate_image", {"size": "invalid"}),
+ ("generate_guided_image", {"control_mode": "invalid", "input_image": "invalid"}),
+ ("generate_guided_image", {"control_mode": "depth", "input_image": "invalid"}),
+ ("inpaint_image", {"input_image": "invalid", "mask_image": "invalid"}),
+ ],
+)
+async def test_invalid_requests_do_not_load(app_config, lazy_backend, tool, arguments) -> None:
+ backend, loads, _ = lazy_backend
+ server = create_server(ImageGenerationService(app_config, backend))
+ async with Client(server) as client:
+ result = await client.call_tool(tool, {"positive_prompt": "cat", **arguments})
+ assert result.is_error is True
+ assert loads == []
+
+
+async def test_invalid_mask_does_not_load(app_config, lazy_backend, image_base64) -> None:
+ backend, loads, _ = lazy_backend
+ async with Client(create_server(ImageGenerationService(app_config, backend))) as client:
+ result = await client.call_tool(
+ "inpaint_image",
+ {"positive_prompt": "cat", "input_image": image_base64, "mask_image": "invalid"},
+ )
+ assert result.is_error is True
+ assert loads == []
+
+
+@pytest.mark.parametrize("tool", ["generate_guided_image", "inpaint_image"])
+async def test_disabled_controlnet_does_not_load(app_config, lazy_backend, image_base64, tool):
+ backend, loads, _ = lazy_backend
+ config = replace(
+ app_config,
+ controlnet=replace(
+ app_config.controlnet,
+ depth=replace(app_config.controlnet.depth, enabled=False),
+ inpaint=replace(app_config.controlnet.inpaint, enabled=False),
+ ),
+ )
+ arguments = {"positive_prompt": "cat", "input_image": image_base64}
+ if tool == "generate_guided_image":
+ arguments["control_mode"] = "depth"
+ else:
+ arguments["mask_image"] = image_base64
+ async with Client(create_server(ImageGenerationService(config, backend))) as client:
+ result = await client.call_tool(tool, arguments)
+ assert result.is_error is True
+ assert loads == []
+
+
+async def test_load_failure_is_tool_error_and_next_call_retries(
+ app_config, lazy_backend, monkeypatch, caplog
+):
+ backend, loads, _ = lazy_backend
+ create_base = backend.factory.create_base
+ attempts = 0
+
+ def fail_once():
+ nonlocal attempts
+ attempts += 1
+ if attempts == 1:
+ raise ConfigError("test model could not load")
+ return create_base()
+
+ monkeypatch.setattr(backend.factory, "create_base", fail_once)
+ async with Client(create_server(ImageGenerationService(app_config, backend))) as client:
+ result = await client.call_tool("generate_image", {"positive_prompt": "cat"})
+ assert result.is_error is True
+ assert "test model could not load" in caplog.text
+ assert backend.base is None
+ assert len((await client.list_tools()).tools) == 3
+ result = await client.call_tool("generate_image", {"positive_prompt": "cat"})
+ assert result.is_error is False
+ assert attempts == 2
+ assert len(loads) == 1
+
+
+async def test_cli_stdio_discovery_forbids_inference_runtime(tmp_path: Path) -> None:
+ checkpoint = tmp_path / "checkpoints" / "placeholder.safetensors"
+ checkpoint.parent.mkdir()
+ checkpoint.touch()
+ config = tmp_path / "config.toml"
+ config.write_text('[model]\nroot = "."\ncheckpoint = "placeholder.safetensors"\n')
+ # Exercise the real CLI while making any inference import or model load fatal.
+ bootstrap = """
+import importlib.abc
+import sys
+class NoInference(importlib.abc.MetaPathFinder):
+ def find_spec(self, fullname, path=None, target=None):
+ if fullname.split('.')[0] in {'torch', 'diffusers', 'controlnet_aux', 'sd_embed'}:
+ raise AssertionError('inference imported during discovery: ' + fullname)
+sys.meta_path.insert(0, NoInference())
+from mcp_sd15.pipeline_factory import PipelineFactory
+def forbidden(self):
+ raise AssertionError('model loaded during discovery')
+PipelineFactory.create_base = forbidden
+from mcp_sd15.cli import main
+raise SystemExit(main())
+"""
+ params = StdioServerParameters(
+ command=sys.executable,
+ args=["-c", bootstrap, "--config", str(config)],
+ )
+ async with (
+ stdio_client(params) as (read, write),
+ ClientSession(read, write, read_timeout_seconds=15) as client,
+ ):
+ await client.initialize()
+ tools = (await client.list_tools()).tools
+ assert {tool.name for tool in tools} == {
+ "generate_image", "generate_guided_image", "inpaint_image"
+ }
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-19 14:26:05 +0000