diff options
| author | tanaka takashi <isazaka@gmail.com> | 2026-09-12 18:04:56 +0900 |
|---|---|---|
| committer | tanaka takashi <isazaka@gmail.com> | 2026-09-12 18:04:56 +0900 |
| commit | cba15a62829a0fe57a03c8031f981e8f790e2efc (patch) | |
| tree | 3bc0ef1a243f356f55c2fa97f63c3bfd5c77e8fc /tests | |
feat: Stable Diffusion 1.x MCPサーバを実装
仕様書・設計書に基づき、ローカルCUDA GPU上のStable Diffusion 1.x互換モデルをMCP経由で利用するサーバを新規実装。
実装内容:
- generate_image、generate_guided_image、inpaint_imageの3ツールをMCP SDK v2で公開
- TOMLからcheckpoint、VAE、複数LoRAと重み、ControlNet、生成設定を読み込み
- モデル種別ごとの固定ディレクトリ解決と起動時の設定・ファイル存在検証を追加
- Depth/MiDaS、OpenPose body/full、ControlNet Inpaintの前処理とパイプライン切替を実装
- FP16/CUDA、1枚固定生成、ControlNetオンデマンドロード、単一GPU直列実行を実装
- Base64 PNG/JPEG入力の形式・サイズ検証と任意のdata URI入力に対応
- seed、サイズ、モード、時刻、処理時間をPNG metadataへ記録
- 任意のローカルPNG保存、生成エラー分類、CUDA OOM後のキャッシュ回復を追加
- --configと--versionを備えたstdio CLIエントリポイントを追加
開発・運用:
- Python 3.11およびuvによる依存・lock管理を構成
- config.example.tomlとセットアップ・運用手順をREADMEへ追加
- Ruff、Pyright、pytestの設定を追加
- 設定、パス解決、schema、画像検証、サービスルーティング、排他制御、MCP画像返却をテスト
検証結果:
- pytest: 30 passed
- Ruff lint/format: passed
- Pyright: 0 errors, 0 warnings
- CLI versionおよびMCPクライアント経由のPNG返却を確認
- 実GPU smoke testはCUDA GPUとモデルファイルがないため対象外
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/__init__.py | 1 | ||||
| -rw-r--r-- | tests/conftest.py | 38 | ||||
| -rw-r--r-- | tests/test_config.py | 122 | ||||
| -rw-r--r-- | tests/test_images.py | 34 | ||||
| -rw-r--r-- | tests/test_schema.py | 43 | ||||
| -rw-r--r-- | tests/test_server.py | 37 | ||||
| -rw-r--r-- | tests/test_service.py | 118 |
7 files changed, 393 insertions, 0 deletions
diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..38bb211 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Test package.""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..d364dbd --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import base64 +import io +from pathlib import Path + +import pytest +from PIL import Image + +from mcp_sd15.config import ( + AppConfig, + ControlModelConfig, + ControlNetConfig, + GenerationConfig, + ModelConfig, + OutputConfig, +) + + +@pytest.fixture +def app_config(tmp_path: Path) -> AppConfig: + return AppConfig( + model=ModelConfig(tmp_path, "base.safetensors", None, ()), + generation=GenerationConfig(), + controlnet=ControlNetConfig( + depth=ControlModelConfig(True, "depth.safetensors", "midas"), + openpose=ControlModelConfig(True, "pose.safetensors"), + inpaint=ControlModelConfig(True, "inpaint.safetensors"), + ), + output=OutputConfig(tmp_path / "outputs", False), + ) + + +@pytest.fixture +def image_base64() -> str: + output = io.BytesIO() + Image.new("RGB", (8, 8), "blue").save(output, format="PNG") + return base64.b64encode(output.getvalue()).decode("ascii") diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..5e4341a --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from mcp_sd15.config import ConfigError, load_config + + +def make_models(root: Path) -> None: + for folder, names in { + "checkpoints": ["base.safetensors"], + "vae": ["vae.safetensors"], + "lora": ["a.safetensors", "b.safetensors"], + "controlnet": ["depth.safetensors", "pose.safetensors", "inpaint.safetensors"], + }.items(): + destination = root / folder + destination.mkdir(parents=True) + for name in names: + (destination / name).touch() + + +def valid_toml(root: Path) -> str: + return f'''[model] +root = "{root.as_posix()}" +checkpoint = "base.safetensors" +vae = "vae.safetensors" +[[model.lora]] +file = "a.safetensors" +weight = 0.8 +[[model.lora]] +file = "b.safetensors" +weight = -0.2 +[generation] +steps = 20 +cfg_scale = 6.5 +scheduler = "euler_a" +default_size = "portrait_512" +[controlnet.depth] +enabled = true +model = "depth.safetensors" +preprocessor = "midas" +[controlnet.openpose] +enabled = true +model = "pose.safetensors" +[controlnet.inpaint] +enabled = true +model = "inpaint.safetensors" +[output] +enabled = true +directory = "generated" +''' + + +def test_loads_and_resolves_all_model_paths(tmp_path: Path) -> None: + root = tmp_path / "models" + make_models(root) + path = tmp_path / "config.toml" + path.write_text(valid_toml(root), encoding="utf-8") + + config = load_config(path) + + assert config.model.checkpoint_path() == root / "checkpoints" / "base.safetensors" + assert config.model.vae_path() == root / "vae" / "vae.safetensors" + assert config.model.lora_paths() == ( + root / "lora" / "a.safetensors", + root / "lora" / "b.safetensors", + ) + assert config.generation.default_size == "portrait_512" + assert config.output.directory == tmp_path / "generated" + assert config.controlnet_path("openpose_full") == root / "controlnet" / "pose.safetensors" + + +@pytest.mark.parametrize( + "line, message", + [ + ('checkpoint = "../base.safetensors"', "file name"), + ("steps = 0", "steps"), + ('scheduler = "unknown"', "scheduler"), + ('default_size = "1024"', "default_size"), + ('preprocessor = "zoe"', "preprocessor"), + ], +) +def test_rejects_invalid_values(tmp_path: Path, line: str, message: str) -> None: + root = tmp_path / "models" + make_models(root) + text = valid_toml(root) + if line.startswith("checkpoint"): + text = text.replace('checkpoint = "base.safetensors"', line) + elif line.startswith("steps"): + text = text.replace("steps = 20", line) + elif line.startswith("scheduler"): + text = text.replace('scheduler = "euler_a"', line) + elif line.startswith("default_size"): + text = text.replace('default_size = "portrait_512"', line) + else: + text = text.replace('preprocessor = "midas"', line) + path = tmp_path / "bad.toml" + path.write_text(text, encoding="utf-8") + + with pytest.raises(ConfigError, match=message): + load_config(path) + + +def test_missing_model_fails_startup_validation(tmp_path: Path) -> None: + root = tmp_path / "models" + make_models(root) + (root / "controlnet" / "depth.safetensors").unlink() + path = tmp_path / "config.toml" + path.write_text(valid_toml(root), encoding="utf-8") + + with pytest.raises(ConfigError, match=r"depth\.safetensors"): + load_config(path) + + +def test_can_skip_path_checks_for_schema_only_use(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + path.write_text( + '[model]\nroot = "missing"\ncheckpoint = "base.safetensors"\n', encoding="utf-8" + ) + config = load_config(path, validate_paths=False) + assert config.model.root == (tmp_path / "missing").resolve() diff --git a/tests/test_images.py b/tests/test_images.py new file mode 100644 index 0000000..f2ac44e --- /dev/null +++ b/tests/test_images.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import base64 +import io + +import pytest +from PIL import Image + +from mcp_sd15.errors import InvalidImageError +from mcp_sd15.images import decode_image, encode_png + + +def test_decodes_raw_base64_and_data_uri(image_base64: str) -> None: + assert decode_image(image_base64).mode == "RGB" + assert decode_image(f"data:image/png;base64,{image_base64}", mask=True).mode == "L" + + +@pytest.mark.parametrize("value", ["not base64!", "data:text/plain;base64,YQ==", ""]) +def test_rejects_invalid_image_input(value: str) -> None: + with pytest.raises(InvalidImageError): + decode_image(value) + + +def test_rejects_non_image_base64() -> None: + value = base64.b64encode(b"plain text").decode() + with pytest.raises(InvalidImageError): + decode_image(value) + + +def test_png_output_contains_metadata() -> None: + data = encode_png(Image.new("RGB", (2, 2)), {"seed": 42, "mode": "test"}) + with Image.open(io.BytesIO(data)) as result: + assert result.format == "PNG" + assert result.info["seed"] == "42" diff --git a/tests/test_schema.py b/tests/test_schema.py new file mode 100644 index 0000000..15a5f9b --- /dev/null +++ b/tests/test_schema.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import pytest + +from mcp_sd15.errors import InvalidRequestError +from mcp_sd15.schema import ( + ControlMode, + GuidedImageRequest, + ImageSize, + InpaintImageRequest, + TextToImageRequest, +) + + +def test_text_request_applies_defaults() -> None: + request = TextToImageRequest.create(" a castle ", default_size="landscape_512") + assert request.positive_prompt == "a castle" + assert request.negative_prompt == "" + assert request.size is ImageSize.LANDSCAPE_512 + assert request.size.dimensions == (512, 384) + + +@pytest.mark.parametrize("prompt", ["", " ", None, 3]) +def test_positive_prompt_is_required(prompt: object) -> None: + with pytest.raises(InvalidRequestError, match="positive_prompt"): + TextToImageRequest.create(prompt) # type: ignore[arg-type] + + +def test_guided_mode_is_an_enum() -> None: + request = GuidedImageRequest.create("cat", "openpose_full", "abc") + assert request.control_mode is ControlMode.OPENPOSE_FULL + with pytest.raises(InvalidRequestError, match="control_mode"): + GuidedImageRequest.create("cat", "canny", "abc") + + +def test_inpaint_requires_both_images() -> None: + with pytest.raises(InvalidRequestError, match="mask_image"): + InpaintImageRequest.create("repair", "abc", "") + + +def test_arbitrary_dimensions_are_not_accepted() -> None: + with pytest.raises(InvalidRequestError, match="size"): + TextToImageRequest.create("cat", size="640x480") diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..4590049 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from mcp import Client +from mcp.types import ImageContent + +from mcp_sd15.server import create_server +from mcp_sd15.service import ImageGenerationService + +from .test_service import FakeBackend + + +async def test_exposes_exactly_three_tools_with_closed_enums(app_config) -> None: + server = create_server(ImageGenerationService(app_config, FakeBackend())) + async with Client(server) as client: + tools = (await client.list_tools()).tools + assert {tool.name for tool in tools} == { + "generate_image", + "generate_guided_image", + "inpaint_image", + } + guided = next(tool for tool in tools if tool.name == "generate_guided_image") + assert guided.input_schema["properties"]["control_mode"]["enum"] == [ + "depth", + "openpose", + "openpose_full", + ] + assert "width" not in guided.input_schema["properties"] + + +async def test_mcp_call_returns_png_image_content(app_config) -> None: + server = create_server(ImageGenerationService(app_config, FakeBackend())) + async with Client(server) as client: + result = await client.call_tool("generate_image", {"positive_prompt": "a cat"}) + assert result.is_error is False + assert len(result.content) == 1 + assert isinstance(result.content[0], ImageContent) + assert result.content[0].mime_type == "image/png" diff --git a/tests/test_service.py b/tests/test_service.py new file mode 100644 index 0000000..da82204 --- /dev/null +++ b/tests/test_service.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import asyncio +import io +import threading +import time +from dataclasses import replace + +import pytest +from PIL import Image + +from mcp_sd15.config import ControlModelConfig, ControlNetConfig +from mcp_sd15.errors import GenerationError, InvalidRequestError +from mcp_sd15.schema import GuidedImageRequest, InpaintImageRequest, TextToImageRequest +from mcp_sd15.service import ImageGenerationService + + +class FakeBackend: + def __init__(self, delay: float = 0) -> None: + self.calls: list[str] = [] + self.delay = delay + self.active = 0 + self.max_active = 0 + self.guard = threading.Lock() + + def _result(self, name: str) -> Image.Image: + with self.guard: + self.active += 1 + self.max_active = max(self.max_active, self.active) + time.sleep(self.delay) + with self.guard: + self.active -= 1 + self.calls.append(name) + return Image.new("RGB", (4, 4), "red") + + def text_to_image(self, request: object, seed: int) -> Image.Image: + return self._result("text") + + def guided_image(self, request: object, image: Image.Image, seed: int) -> Image.Image: + assert image.mode == "RGB" + return self._result("guided") + + def inpaint_image( + self, request: object, image: Image.Image, mask: Image.Image, seed: int + ) -> Image.Image: + assert mask.mode == "L" + return self._result("inpaint") + + +async def test_routes_all_generation_types(app_config, image_base64: str) -> None: + backend = FakeBackend() + service = ImageGenerationService(app_config, backend) + + text = await service.generate_image(TextToImageRequest.create("cat")) + guided = await service.generate_guided_image( + GuidedImageRequest.create("cat", "depth", image_base64) + ) + inpaint = await service.inpaint_image( + InpaintImageRequest.create("cat", image_base64, image_base64) + ) + + assert backend.calls == ["text", "guided", "inpaint"] + for result in (text, guided, inpaint): + assert result.png_data.startswith(b"\x89PNG") + assert isinstance(result.metadata["seed"], int) + + +async def test_generation_is_strictly_serialized(app_config) -> None: + backend = FakeBackend(delay=0.03) + service = ImageGenerationService(app_config, backend) + await asyncio.gather( + *(service.generate_image(TextToImageRequest.create(f"image {i}")) for i in range(4)) + ) + assert backend.max_active == 1 + + +async def test_disabled_controlnet_fails_before_backend(app_config, image_base64: str) -> None: + disabled = replace( + app_config, + controlnet=ControlNetConfig( + depth=ControlModelConfig(False, None, "midas"), + openpose=app_config.controlnet.openpose, + inpaint=app_config.controlnet.inpaint, + ), + ) + service = ImageGenerationService(disabled, FakeBackend()) + with pytest.raises(InvalidRequestError, match="disabled"): + await service.generate_guided_image(GuidedImageRequest.create("cat", "depth", image_base64)) + + +async def test_backend_error_is_sanitized(app_config) -> None: + class BrokenBackend(FakeBackend): + def text_to_image(self, request: object, seed: int) -> Image.Image: + raise RuntimeError("boom") + + service = ImageGenerationService(app_config, BrokenBackend()) + with pytest.raises(GenerationError, match="boom"): + await service.generate_image(TextToImageRequest.create("cat")) + + +async def test_cuda_oom_is_classified(app_config) -> None: + class OomBackend(FakeBackend): + def text_to_image(self, request: object, seed: int) -> Image.Image: + raise RuntimeError("CUDA out of memory while allocating") + + service = ImageGenerationService(app_config, OomBackend()) + with pytest.raises(GenerationError, match="CUDA out of memory"): + await service.generate_image(TextToImageRequest.create("cat")) + + +async def test_optional_output_is_saved(app_config) -> None: + enabled = replace(app_config, output=replace(app_config.output, enabled=True)) + service = ImageGenerationService(enabled, FakeBackend()) + await service.generate_image(TextToImageRequest.create("cat")) + files = list(enabled.output.directory.glob("*.png")) + assert len(files) == 1 + with Image.open(io.BytesIO(files[0].read_bytes())) as image: + assert image.format == "PNG" |
