summaryrefslogtreecommitdiff
path: root/src/mcp_sd15/service.py
diff options
context:
space:
mode:
authortanaka takashi <isazaka@gmail.com>2026-09-12 18:04:56 +0900
committertanaka takashi <isazaka@gmail.com>2026-09-12 18:04:56 +0900
commitcba15a62829a0fe57a03c8031f981e8f790e2efc (patch)
tree3bc0ef1a243f356f55c2fa97f63c3bfd5c77e8fc /src/mcp_sd15/service.py
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 'src/mcp_sd15/service.py')
-rw-r--r--src/mcp_sd15/service.py125
1 files changed, 125 insertions, 0 deletions
diff --git a/src/mcp_sd15/service.py b/src/mcp_sd15/service.py
new file mode 100644
index 0000000..65b62e7
--- /dev/null
+++ b/src/mcp_sd15/service.py
@@ -0,0 +1,125 @@
+"""Serialized generation service and persistence policy."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import secrets
+import time
+import uuid
+from datetime import UTC, datetime
+from typing import Protocol
+
+from PIL import Image
+
+from .config import AppConfig
+from .errors import GenerationError, InvalidRequestError, SD15Error
+from .images import decode_image, encode_png
+from .schema import (
+ GenerationResult,
+ GuidedImageRequest,
+ InpaintImageRequest,
+ TextToImageRequest,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class GenerationBackend(Protocol):
+ """Synchronous GPU boundary; the service runs it in a worker thread."""
+
+ def text_to_image(self, request: TextToImageRequest, seed: int) -> Image.Image: ...
+
+ def guided_image(
+ self, request: GuidedImageRequest, image: Image.Image, seed: int
+ ) -> Image.Image: ...
+
+ def inpaint_image(
+ self,
+ request: InpaintImageRequest,
+ image: Image.Image,
+ mask: Image.Image,
+ seed: int,
+ ) -> Image.Image: ...
+
+
+class ImageGenerationService:
+ """Validate/decode requests and serialize access to the single GPU."""
+
+ def __init__(self, config: AppConfig, backend: GenerationBackend) -> None:
+ self.config = config
+ self.backend = backend
+ self._generation_lock = asyncio.Lock()
+
+ async def generate_image(self, request: TextToImageRequest) -> GenerationResult:
+ return await self._run("text_to_image", request)
+
+ async def generate_guided_image(self, request: GuidedImageRequest) -> GenerationResult:
+ control = self.config.controlnet.for_mode(request.control_mode)
+ if not control.enabled:
+ raise InvalidRequestError(f"ControlNet mode is disabled: {request.control_mode}")
+ image = decode_image(request.input_image)
+ return await self._run("guided_image", request, image)
+
+ async def inpaint_image(self, request: InpaintImageRequest) -> GenerationResult:
+ if not self.config.controlnet.inpaint.enabled:
+ raise InvalidRequestError("ControlNet mode is disabled: inpaint")
+ image = decode_image(request.input_image)
+ mask = decode_image(request.mask_image, mask=True)
+ return await self._run("inpaint_image", request, image, mask)
+
+ async def _run(
+ self,
+ method: str,
+ request: TextToImageRequest | GuidedImageRequest | InpaintImageRequest,
+ *images: Image.Image,
+ ) -> GenerationResult:
+ seed = secrets.randbelow(2**63)
+ started = time.monotonic()
+ async with self._generation_lock:
+ logger.info("generation started mode=%s size=%s seed=%d", method, request.size, seed)
+ try:
+ function = getattr(self.backend, method)
+ output = await asyncio.to_thread(function, request, *images, seed)
+ if not isinstance(output, Image.Image):
+ raise GenerationError("generation backend returned a non-image result")
+ except SD15Error:
+ raise
+ except Exception as exc:
+ logger.exception("generation failed mode=%s", method)
+ if "out of memory" in str(exc).lower():
+ self._recover_cuda()
+ raise GenerationError("CUDA out of memory") from exc
+ raise GenerationError(f"generation failed: {exc}") from exc
+
+ metadata: dict[str, str | int | float] = {
+ "seed": seed,
+ "size": str(request.size),
+ "mode": method,
+ "timestamp": datetime.now(UTC).isoformat(),
+ "elapsed_seconds": round(time.monotonic() - started, 3),
+ }
+ png_data = encode_png(output, metadata)
+ if self.config.output.enabled:
+ self._save(png_data)
+ logger.info("generation completed mode=%s seed=%d", method, seed)
+ return GenerationResult(png_data, metadata)
+
+ def _save(self, png_data: bytes) -> None:
+ directory = self.config.output.directory
+ try:
+ directory.mkdir(parents=True, exist_ok=True)
+ destination = directory / f"{uuid.uuid4()}.png"
+ destination.write_bytes(png_data)
+ except OSError as exc:
+ raise GenerationError(f"could not save generated image: {exc}") from exc
+
+ @staticmethod
+ def _recover_cuda() -> None:
+ try:
+ import torch
+
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ except ImportError:
+ pass