From 583a56f1ac3d9be9ce2f8c1715d77423347d4a4e Mon Sep 17 00:00:00 2001 From: tanaka takashi Date: Wed, 14 Oct 2026 02:21:50 +0900 Subject: fix: 画像生成モデルのロードを初回生成まで遅延する MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stdio MCPのツール一覧取得でも、受付開始前にModelManagerがSD本体・ 外部VAE・LoRAをロードし、CUDAへ転送していた。画像生成しない通常の 質問でもこの起動コストが発生するため、有効な生成要求の実行時まで モデル初期化を遅らせ、一覧取得時の不要な遅延とVRAM消費を除去する。 実装: - ModelManagerの構築時は未ロード状態とし、冪等なload()を追加する。 - 初期化ロックで二重ロードを防ぎ、成功時だけモデルを保持する。 失敗時は未ロード状態を維持し、次の要求で再試行できるようにする。 - 3種類の生成メソッドからload()を呼び、既存サービスの入力検証・ 生成ロックを経由したワーカースレッド内で初期化する。 - ControlNetと前処理モデルの既存の遅延ロード・モデル共有を維持する。 - 起動・設定読み込みと実際のモデルロードのログを分離し、ロードの 開始・完了・失敗をstderrへ出力する。 - benchmarkの計測区間で明示的にload()を呼び、ロード時間とVRAMの 計測の意味を維持する。READMEにロードタイミングと確認手順を追記する。 変更範囲: - stdioのプロセス起動・終了方式を維持し、常駐化やHTTP切り替え、 チャット側の接続管理・一覧キャッシュ・アイドル解放は追加しない。 - MCPスキーマ、設定形式、CLI引数、プロンプト・seed処理、scheduler、 生成サイズ、PNGメタデータ・保存仕様は変更しない。 - 設定とモデルファイルの存在確認は起動時に残し、モデル互換性・ CUDA・推論依存関係のエラー検出は初回生成時へ移す。 - 同一プロセス内ではモデルを再利用し、終了時に解放する。 別の生成プロセスでの再ロードは引き続き許容する。 回帰テスト: - 修正前の既存69件に対し、修正後は追加18件を含む87件すべて成功。 既存テストと期待値は変更していない。 - CLI/stdioでの一覧取得時にモデルロードと推論ライブラリのimportを 禁止しても成功することを確認する。 - 全生成方式の初回ロード、ワーカースレッド実行、同時要求、冪等性、 不正入力・無効なControlNet指定、失敗後の再試行、ログを検証する。 - ruff check .は成功、pyrightは0 errors / 0 warnings。 実GPU・stdioでの確認: - 現在の設定（外部VAE、2件のLoRA、LCM、6 steps）を維持して計測した。 - 起動から一覧取得までの時間は8.233秒から0.671秒に短縮した。 - 一覧取得前後のGPU全体の使用量は修正前808→3233 MiBに対し、 修正後682→682 MiBで増加しなかった。 - 同じ接続で2枚生成してもベースモデルのロードは1回で、一覧取得では 0回だった。ロード時間は7.812秒で、初回生成要求に移った。 - 同じ設定・プロンプト・seedによる修正前後の画像は画素が完全一致。 PNG形式・寸法・メタデータと、返却PNGおよび保存PNGの一致を確認した。 - 最終版スクリプトでの再確認でも一覧取得0.655秒、初回生成14.089秒、 2枚目1.221秒で、画素・保存結果の一致と生成プロセス終了を確認した。 - 通常質問の親アプリの処理経路でも、LLMのみをスタブ化し、実際の 設定済みuvコマンドで一覧取得・3ツールの登録を確認した。 - 計測スクリプトと詳細な検証報告を追加する。生のJSON・stderr・PNGは gitignore対象のdata/verification-lazy-loadに保存する。 確認の限界: - Windowsのプロセス別VRAMはN/Aのため、GPU全体の使用量・プロセス一覧・ tasklistによる終了確認で評価した。全体値には他プロセスの変動がある。 - 計測値は単一セッションの観測であり、統計的な速度保証ではない。 - ブラウザ操作と実LLMを含むチャット全体、および現在の設定で無効な ControlNetの実GPU生成は未実施。誘導生成・Inpaintのロード順は スタブを用いた自動テストで確認した。 --- README.md | 21 +++- docs/lazy-loading-verification.md | 75 ++++++++++++ script/benchmark_generate_image.py | 1 + script/smoke_stdio.py | 153 +++++++++++++++++++++++ src/mcp_sd15/cli.py | 3 +- src/mcp_sd15/model_manager.py | 36 +++++- tests/test_lazy_loading.py | 245 +++++++++++++++++++++++++++++++++++++ 7 files changed, 530 insertions(+), 4 deletions(-) create mode 100644 docs/lazy-loading-verification.md create mode 100644 script/smoke_stdio.py create mode 100644 tests/test_lazy_loading.py 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`の モデルは `/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" + } -- cgit v1.2.3