From cba15a62829a0fe57a03c8031f981e8f790e2efc Mon Sep 17 00:00:00 2001 From: tanaka takashi Date: Sat, 12 Sep 2026 18:04:56 +0900 Subject: feat: Stable Diffusion 1.x MCPサーバを実装 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 仕様書・設計書に基づき、ローカル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とモデルファイルがないため対象外 --- src/mcp_sd15/service.py | 125 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 src/mcp_sd15/service.py (limited to 'src/mcp_sd15/service.py') 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 -- cgit v1.2.3