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 /src/mcp_sd15 | |
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')
| -rw-r--r-- | src/mcp_sd15/__init__.py | 3 | ||||
| -rw-r--r-- | src/mcp_sd15/__main__.py | 3 | ||||
| -rw-r--r-- | src/mcp_sd15/cli.py | 48 | ||||
| -rw-r--r-- | src/mcp_sd15/config.py | 237 | ||||
| -rw-r--r-- | src/mcp_sd15/errors.py | 21 | ||||
| -rw-r--r-- | src/mcp_sd15/images.py | 62 | ||||
| -rw-r--r-- | src/mcp_sd15/model_manager.py | 142 | ||||
| -rw-r--r-- | src/mcp_sd15/pipeline_factory.py | 133 | ||||
| -rw-r--r-- | src/mcp_sd15/preprocessors/__init__.py | 1 | ||||
| -rw-r--r-- | src/mcp_sd15/preprocessors/depth.py | 24 | ||||
| -rw-r--r-- | src/mcp_sd15/preprocessors/openpose.py | 27 | ||||
| -rw-r--r-- | src/mcp_sd15/schema.py | 142 | ||||
| -rw-r--r-- | src/mcp_sd15/server.py | 66 | ||||
| -rw-r--r-- | src/mcp_sd15/service.py | 125 | ||||
| -rw-r--r-- | src/mcp_sd15/tools/__init__.py | 1 | ||||
| -rw-r--r-- | src/mcp_sd15/tools/generate_guided_image.py | 23 | ||||
| -rw-r--r-- | src/mcp_sd15/tools/generate_image.py | 19 | ||||
| -rw-r--r-- | src/mcp_sd15/tools/inpaint_image.py | 23 |
18 files changed, 1100 insertions, 0 deletions
diff --git a/src/mcp_sd15/__init__.py b/src/mcp_sd15/__init__.py new file mode 100644 index 0000000..0c8e718 --- /dev/null +++ b/src/mcp_sd15/__init__.py @@ -0,0 +1,3 @@ +"""Local Stable Diffusion 1.x MCP server.""" + +__version__ = "0.1.0" diff --git a/src/mcp_sd15/__main__.py b/src/mcp_sd15/__main__.py new file mode 100644 index 0000000..eb53e2f --- /dev/null +++ b/src/mcp_sd15/__main__.py @@ -0,0 +1,3 @@ +from .cli import main + +raise SystemExit(main()) diff --git a/src/mcp_sd15/cli.py b/src/mcp_sd15/cli.py new file mode 100644 index 0000000..8bc1e20 --- /dev/null +++ b/src/mcp_sd15/cli.py @@ -0,0 +1,48 @@ +"""Command-line entry point.""" + +from __future__ import annotations + +import argparse +import logging +from collections.abc import Sequence +from pathlib import Path + +from . import __version__ +from .config import load_config +from .errors import SD15Error +from .model_manager import ModelManager +from .server import create_server +from .service import ImageGenerationService + + +def parser() -> argparse.ArgumentParser: + result = argparse.ArgumentParser(prog="mcp-sd15") + result.add_argument("--config", type=Path, default=Path("config.toml")) + result.add_argument("--version", action="version", version=f"%(prog)s {__version__}") + return result + + +def main(argv: Sequence[str] | None = None) -> int: + args = parser().parse_args(argv) + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + try: + config = load_config(args.config) + logging.getLogger(__name__).info( + "loading checkpoint=%s loras=%d", + config.model.checkpoint_path(), + len(config.model.loras), + ) + backend = ModelManager(config) + service = ImageGenerationService(config, backend) + create_server(service).run() + except SD15Error as exc: + logging.getLogger(__name__).error("startup failed: %s", exc) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/mcp_sd15/config.py b/src/mcp_sd15/config.py new file mode 100644 index 0000000..5bbe225 --- /dev/null +++ b/src/mcp_sd15/config.py @@ -0,0 +1,237 @@ +"""TOML configuration loading and model path confinement.""" + +from __future__ import annotations + +import math +import tomllib +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from .errors import ConfigError + +VALID_SIZES = frozenset({"square_512", "landscape_512", "portrait_512"}) +VALID_SCHEDULERS = frozenset({"dpmpp_2m", "ddim", "euler", "euler_a"}) + + +@dataclass(frozen=True) +class LoRAConfig: + file: str + weight: float = 1.0 + + +@dataclass(frozen=True) +class ModelConfig: + root: Path + checkpoint: str + vae: str | None + loras: tuple[LoRAConfig, ...] + + def checkpoint_path(self) -> Path: + return self.root / "checkpoints" / self.checkpoint + + def vae_path(self) -> Path | None: + return self.root / "vae" / self.vae if self.vae else None + + def lora_paths(self) -> tuple[Path, ...]: + return tuple(self.root / "lora" / item.file for item in self.loras) + + +@dataclass(frozen=True) +class GenerationConfig: + steps: int = 25 + cfg_scale: float = 7.0 + scheduler: str = "dpmpp_2m" + default_size: str = "square_512" + + +@dataclass(frozen=True) +class ControlModelConfig: + enabled: bool + model: str | None + preprocessor: str | None = None + + +@dataclass(frozen=True) +class ControlNetConfig: + depth: ControlModelConfig + openpose: ControlModelConfig + inpaint: ControlModelConfig + + def for_mode(self, mode: str) -> ControlModelConfig: + key = "openpose" if mode == "openpose_full" else mode + if key not in {"depth", "openpose", "inpaint"}: + raise ConfigError(f"unknown ControlNet mode: {mode}") + return getattr(self, key) + + +@dataclass(frozen=True) +class OutputConfig: + directory: Path + enabled: bool = False + + +@dataclass(frozen=True) +class AppConfig: + model: ModelConfig + generation: GenerationConfig + controlnet: ControlNetConfig + output: OutputConfig + + def controlnet_path(self, mode: str) -> Path: + item = self.controlnet.for_mode(mode) + if not item.enabled or not item.model: + raise ConfigError(f"ControlNet mode is disabled: {mode}") + return self.model.root / "controlnet" / item.model + + +def _table(data: dict[str, Any], key: str, *, required: bool = False) -> dict[str, Any]: + value = data.get(key) + if value is None and not required: + return {} + if not isinstance(value, dict): + raise ConfigError(f"[{key}] must be a TOML table") + return value + + +def _string(table: dict[str, Any], key: str, *, required: bool = False) -> str | None: + value = table.get(key) + if value is None and not required: + return None + if not isinstance(value, str) or not value.strip(): + raise ConfigError(f"{key} must be a non-empty string") + return value.strip() + + +def _filename(value: str | None, label: str) -> str | None: + if value is None: + return None + path = Path(value) + if path.name != value or path.is_absolute() or value in {".", ".."}: + raise ConfigError(f"{label} must be a file name, not a path") + return value + + +def _bool(table: dict[str, Any], key: str, default: bool) -> bool: + value = table.get(key, default) + if not isinstance(value, bool): + raise ConfigError(f"{key} must be a boolean") + return value + + +def _number(table: dict[str, Any], key: str, default: float) -> float: + value = table.get(key, default) + if isinstance(value, bool) or not isinstance(value, int | float): + raise ConfigError(f"{key} must be a number") + result = float(value) + if not math.isfinite(result): + raise ConfigError(f"{key} must be finite") + return result + + +def _control(table: dict[str, Any], name: str) -> ControlModelConfig: + enabled = _bool(table, "enabled", False) + model = _filename(_string(table, "model"), f"controlnet.{name}.model") + if enabled and model is None: + raise ConfigError(f"controlnet.{name}.model is required when enabled") + preprocessor = _string(table, "preprocessor") + if name == "depth" and preprocessor not in {None, "midas"}: + raise ConfigError("controlnet.depth.preprocessor must be 'midas'") + return ControlModelConfig( + enabled, model, preprocessor or ("midas" if name == "depth" else None) + ) + + +def load_config(path: str | Path, *, validate_paths: bool = True) -> AppConfig: + """Load a configuration file, validating all values and local model paths.""" + config_path = Path(path).expanduser().resolve() + try: + with config_path.open("rb") as stream: + raw = tomllib.load(stream) + except (OSError, tomllib.TOMLDecodeError) as exc: + raise ConfigError(f"cannot read config {config_path}: {exc}") from exc + + model_raw = _table(raw, "model", required=True) + root_value = _string(model_raw, "root", required=True) + assert root_value is not None + root = Path(root_value).expanduser() + if not root.is_absolute(): + root = config_path.parent / root + root = root.resolve() + checkpoint = _filename(_string(model_raw, "checkpoint", required=True), "model.checkpoint") + assert checkpoint is not None + vae = _filename(_string(model_raw, "vae"), "model.vae") + + lora_raw = model_raw.get("lora", []) + if not isinstance(lora_raw, list): + raise ConfigError("model.lora must be an array of tables") + loras: list[LoRAConfig] = [] + for index, value in enumerate(lora_raw): + if not isinstance(value, dict): + raise ConfigError(f"model.lora[{index}] must be a table") + file = _filename(_string(value, "file", required=True), f"model.lora[{index}].file") + weight = _number(value, "weight", 1.0) + if not -4.0 <= weight <= 4.0: + raise ConfigError(f"model.lora[{index}].weight must be between -4 and 4") + assert file is not None + loras.append(LoRAConfig(file, weight)) + if len({item.file for item in loras}) != len(loras): + raise ConfigError("LoRA file names must be unique") + + generation_raw = _table(raw, "generation") + steps_value = generation_raw.get("steps", 25) + if ( + isinstance(steps_value, bool) + or not isinstance(steps_value, int) + or not 1 <= steps_value <= 150 + ): + raise ConfigError("generation.steps must be an integer between 1 and 150") + cfg_scale = _number(generation_raw, "cfg_scale", 7.0) + if not 0 <= cfg_scale <= 30: + raise ConfigError("generation.cfg_scale must be between 0 and 30") + scheduler = _string(generation_raw, "scheduler") or "dpmpp_2m" + if scheduler not in VALID_SCHEDULERS: + raise ConfigError(f"unsupported scheduler: {scheduler}") + default_size = _string(generation_raw, "default_size") or "square_512" + if default_size not in VALID_SIZES: + raise ConfigError(f"unsupported default_size: {default_size}") + + control_raw = _table(raw, "controlnet") + controlnet = ControlNetConfig( + depth=_control(_table(control_raw, "depth"), "depth"), + openpose=_control(_table(control_raw, "openpose"), "openpose"), + inpaint=_control(_table(control_raw, "inpaint"), "inpaint"), + ) + output_raw = _table(raw, "output") + output_value = _string(output_raw, "directory") or "./outputs" + output_dir = Path(output_value).expanduser() + if not output_dir.is_absolute(): + output_dir = config_path.parent / output_dir + + config = AppConfig( + model=ModelConfig(root, checkpoint, vae, tuple(loras)), + generation=GenerationConfig(steps_value, cfg_scale, scheduler, default_size), + controlnet=controlnet, + output=OutputConfig(output_dir.resolve(), _bool(output_raw, "enabled", False)), + ) + if validate_paths: + validate_model_paths(config) + return config + + +def validate_model_paths(config: AppConfig) -> None: + """Fail before GPU initialization when a configured model file is absent.""" + if not config.model.root.is_dir(): + raise ConfigError(f"model root directory not found: {config.model.root}") + required = [config.model.checkpoint_path()] + vae = config.model.vae_path() + if vae: + required.append(vae) + required.extend(config.model.lora_paths()) + for mode in ("depth", "openpose", "inpaint"): + item = config.controlnet.for_mode(mode) + if item.enabled: + required.append(config.controlnet_path(mode)) + missing = [str(path) for path in required if not path.is_file()] + if missing: + raise ConfigError("model file(s) not found: " + ", ".join(missing)) diff --git a/src/mcp_sd15/errors.py b/src/mcp_sd15/errors.py new file mode 100644 index 0000000..59ba567 --- /dev/null +++ b/src/mcp_sd15/errors.py @@ -0,0 +1,21 @@ +"""Application error types exposed at architectural boundaries.""" + + +class SD15Error(Exception): + """Base application error.""" + + +class ConfigError(SD15Error): + """Configuration is invalid or incomplete.""" + + +class InvalidRequestError(SD15Error): + """A tool request is invalid.""" + + +class InvalidImageError(InvalidRequestError): + """An input is not a supported image.""" + + +class GenerationError(SD15Error): + """The inference backend failed.""" diff --git a/src/mcp_sd15/images.py b/src/mcp_sd15/images.py new file mode 100644 index 0000000..4b9cf3c --- /dev/null +++ b/src/mcp_sd15/images.py @@ -0,0 +1,62 @@ +"""Safe conversion between MCP Base64 inputs and Pillow images.""" + +from __future__ import annotations + +import base64 +import binascii +import io +import json +from typing import Any + +from PIL import Image, PngImagePlugin, UnidentifiedImageError + +from .errors import InvalidImageError + +MAX_ENCODED_LENGTH = 36_000_000 +MAX_PIXELS = 40_000_000 +SUPPORTED_FORMATS = frozenset({"PNG", "JPEG"}) + + +def decode_image(value: str, *, mask: bool = False) -> Image.Image: + """Decode a bounded PNG/JPEG Base64 value or data URI and fully validate it.""" + if not isinstance(value, str) or len(value) > MAX_ENCODED_LENGTH: + raise InvalidImageError("image Base64 input is missing or too large") + payload = value.strip() + if payload.startswith("data:"): + try: + header, payload = payload.split(",", 1) + except ValueError as exc: + raise InvalidImageError("malformed image data URI") from exc + if header.lower() not in { + "data:image/png;base64", + "data:image/jpeg;base64", + "data:image/jpg;base64", + }: + raise InvalidImageError("data URI must contain a Base64 PNG or JPEG") + try: + raw = base64.b64decode(payload, validate=True) + except (binascii.Error, ValueError) as exc: + raise InvalidImageError("image is not valid Base64") from exc + try: + with Image.open(io.BytesIO(raw)) as opened: + if opened.format not in SUPPORTED_FORMATS: + raise InvalidImageError("only PNG and JPEG images are supported") + if opened.width * opened.height > MAX_PIXELS: + raise InvalidImageError("image pixel count exceeds the safety limit") + opened.load() + return opened.convert("L" if mask else "RGB") + except InvalidImageError: + raise + except (OSError, UnidentifiedImageError, Image.DecompressionBombError) as exc: + raise InvalidImageError("input is not a valid PNG or JPEG image") from exc + + +def encode_png(image: Image.Image, metadata: dict[str, Any]) -> bytes: + """Serialize one RGB image as PNG with stable textual generation metadata.""" + info = PngImagePlugin.PngInfo() + for key, value in metadata.items(): + text = json.dumps(value) if isinstance(value, (dict, list)) else str(value) + info.add_text(str(key), text) + output = io.BytesIO() + image.convert("RGB").save(output, format="PNG", pnginfo=info) + return output.getvalue() diff --git a/src/mcp_sd15/model_manager.py b/src/mcp_sd15/model_manager.py new file mode 100644 index 0000000..c2bd166 --- /dev/null +++ b/src/mcp_sd15/model_manager.py @@ -0,0 +1,142 @@ +"""Diffusers-backed generation and on-demand ControlNet management.""" + +from __future__ import annotations + +import gc +from typing import Any + +import numpy as np +from PIL import Image + +from .config import AppConfig +from .errors import GenerationError +from .pipeline_factory import PipelineFactory +from .preprocessors.depth import DepthPreprocessor +from .preprocessors.openpose import OpenPosePreprocessor +from .schema import GuidedImageRequest, InpaintImageRequest, TextToImageRequest + + +class ModelManager: + """Own all model state; construction validates model compatibility at startup.""" + + def __init__(self, config: AppConfig) -> None: + self.config = config + self.factory = PipelineFactory(config) + self.base = self.factory.create_base() + self._control_mode: str | None = None + self._controlnet: Any | None = None + self._control_pipeline: Any | None = None + self._depth = DepthPreprocessor() + self._openpose = OpenPosePreprocessor() + + def text_to_image(self, request: TextToImageRequest, seed: int) -> Image.Image: + width, height = request.size.dimensions + result = self.base( + prompt=request.positive_prompt, + negative_prompt=request.negative_prompt, + width=width, + height=height, + num_inference_steps=self.config.generation.steps, + guidance_scale=self.config.generation.cfg_scale, + num_images_per_prompt=1, + generator=self._generator(seed), + ) + return self._one_image(result) + + def guided_image( + self, request: GuidedImageRequest, image: Image.Image, seed: int + ) -> Image.Image: + width, height = request.size.dimensions + image = image.resize((width, height), Image.Resampling.LANCZOS) + if request.control_mode == "depth": + control_image = self._depth(image) + else: + control_image = self._openpose(image, full=request.control_mode == "openpose_full") + pipeline = self._pipeline_for(str(request.control_mode), inpaint=False) + result = pipeline( + prompt=request.positive_prompt, + negative_prompt=request.negative_prompt, + image=control_image, + width=width, + height=height, + num_inference_steps=self.config.generation.steps, + guidance_scale=self.config.generation.cfg_scale, + num_images_per_prompt=1, + generator=self._generator(seed), + ) + return self._one_image(result) + + def inpaint_image( + self, + request: InpaintImageRequest, + image: Image.Image, + mask: Image.Image, + seed: int, + ) -> Image.Image: + width, height = request.size.dimensions + image = image.resize((width, height), Image.Resampling.LANCZOS) + mask = mask.resize((width, height), Image.Resampling.NEAREST) + control_image = self._inpaint_condition(image, mask) + pipeline = self._pipeline_for("inpaint", inpaint=True) + result = pipeline( + prompt=request.positive_prompt, + negative_prompt=request.negative_prompt, + image=image, + mask_image=mask, + control_image=control_image, + width=width, + height=height, + num_inference_steps=self.config.generation.steps, + guidance_scale=self.config.generation.cfg_scale, + num_images_per_prompt=1, + generator=self._generator(seed), + ) + return self._one_image(result) + + def _pipeline_for(self, mode: str, *, inpaint: bool) -> Any: + model_mode = "openpose" if mode == "openpose_full" else mode + if self._control_mode == model_mode and self._control_pipeline is not None: + return self._control_pipeline + self._release_controlnet() + self._controlnet = self.factory.create_controlnet(self.config.controlnet_path(model_mode)) + self._control_pipeline = ( + self.factory.create_inpaint_pipeline(self.base, self._controlnet) + if inpaint + else self.factory.create_guided_pipeline(self.base, self._controlnet) + ) + self._control_mode = model_mode + return self._control_pipeline + + def _release_controlnet(self) -> None: + self._control_pipeline = None + self._controlnet = None + self._control_mode = None + gc.collect() + try: + import torch + + torch.cuda.empty_cache() + except ImportError: + pass + + @staticmethod + def _generator(seed: int) -> Any: + import torch + + return torch.Generator(device="cuda").manual_seed(seed) + + @staticmethod + def _inpaint_condition(image: Image.Image, mask: Image.Image) -> Any: + import torch + + image_array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 + mask_array = np.asarray(mask.convert("L"), dtype=np.float32) / 255.0 + image_array[mask_array > 0.5] = -1.0 + return torch.from_numpy(image_array[None].transpose(0, 3, 1, 2)) + + @staticmethod + def _one_image(result: Any) -> Image.Image: + images = getattr(result, "images", None) + if not images or len(images) != 1 or not isinstance(images[0], Image.Image): + raise GenerationError("Diffusers did not return exactly one image") + return images[0] diff --git a/src/mcp_sd15/pipeline_factory.py b/src/mcp_sd15/pipeline_factory.py new file mode 100644 index 0000000..cd92690 --- /dev/null +++ b/src/mcp_sd15/pipeline_factory.py @@ -0,0 +1,133 @@ +"""Lazy Diffusers construction kept outside the MCP layer.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from .config import AppConfig +from .errors import ConfigError + + +class PipelineFactory: + """Build SD1.x pipelines in FP16 without network access.""" + + def __init__(self, config: AppConfig) -> None: + self.config = config + + @staticmethod + def _runtime() -> tuple[Any, Any]: + try: + import diffusers + import torch + except ImportError as exc: + raise ConfigError(f"inference dependency is unavailable: {exc}") from exc + if not torch.cuda.is_available(): + raise ConfigError("CUDA GPU is required but torch.cuda.is_available() is false") + return torch, diffusers + + def create_base(self) -> Any: + torch, diffusers = self._runtime() + try: + pipeline = diffusers.StableDiffusionPipeline.from_single_file( + str(self.config.model.checkpoint_path()), + torch_dtype=torch.float16, + local_files_only=True, + ) + vae_path = self.config.model.vae_path() + if vae_path is not None: + pipeline.vae = diffusers.AutoencoderKL.from_single_file( + str(vae_path), + torch_dtype=torch.float16, + local_files_only=True, + ) + adapter_names: list[str] = [] + adapter_weights: list[float] = [] + for index, (item, path) in enumerate( + zip(self.config.model.loras, self.config.model.lora_paths(), strict=True) + ): + name = f"configured_lora_{index}" + pipeline.load_lora_weights( + str(path.parent), + weight_name=path.name, + adapter_name=name, + local_files_only=True, + ) + adapter_names.append(name) + adapter_weights.append(item.weight) + if adapter_names: + pipeline.set_adapters(adapter_names, adapter_weights=adapter_weights) + pipeline.scheduler = self._scheduler(pipeline.scheduler, diffusers) + pipeline.set_progress_bar_config(disable=True) + return pipeline.to("cuda") + except ConfigError: + raise + except Exception as exc: + raise ConfigError(f"failed to load SD1.x model configuration: {exc}") from exc + + def create_controlnet(self, path: Path) -> Any: + torch, diffusers = self._runtime() + try: + return diffusers.ControlNetModel.from_single_file( + str(path), torch_dtype=torch.float16, local_files_only=True + ).to("cuda") + except Exception as exc: + raise ConfigError(f"failed to load ControlNet model {path}: {exc}") from exc + + @staticmethod + def create_guided_pipeline(base: Any, controlnet: Any) -> Any: + from diffusers.pipelines.controlnet.pipeline_controlnet import ( + StableDiffusionControlNetPipeline, + ) + + components: dict[str, Any] = dict( + vae=base.vae, + text_encoder=base.text_encoder, + tokenizer=base.tokenizer, + unet=base.unet, + controlnet=controlnet, + scheduler=base.scheduler, + safety_checker=base.safety_checker, + feature_extractor=base.feature_extractor, + requires_safety_checker=getattr(base, "requires_safety_checker", False), + ) + if getattr(base, "image_encoder", None) is not None: + components["image_encoder"] = base.image_encoder + pipeline = StableDiffusionControlNetPipeline(**components) + pipeline.set_progress_bar_config(disable=True) + return pipeline.to("cuda") + + @staticmethod + def create_inpaint_pipeline(base: Any, controlnet: Any) -> Any: + from diffusers.pipelines.controlnet.pipeline_controlnet_inpaint import ( + StableDiffusionControlNetInpaintPipeline, + ) + + components: dict[str, Any] = dict( + vae=base.vae, + text_encoder=base.text_encoder, + tokenizer=base.tokenizer, + unet=base.unet, + controlnet=controlnet, + scheduler=base.scheduler, + safety_checker=base.safety_checker, + feature_extractor=base.feature_extractor, + requires_safety_checker=getattr(base, "requires_safety_checker", False), + ) + if getattr(base, "image_encoder", None) is not None: + components["image_encoder"] = base.image_encoder + pipeline = StableDiffusionControlNetInpaintPipeline(**components) + pipeline.set_progress_bar_config(disable=True) + return pipeline.to("cuda") + + def _scheduler(self, scheduler: Any, diffusers: Any) -> Any: + classes = { + "dpmpp_2m": diffusers.DPMSolverMultistepScheduler, + "ddim": diffusers.DDIMScheduler, + "euler": diffusers.EulerDiscreteScheduler, + "euler_a": diffusers.EulerAncestralDiscreteScheduler, + } + options = dict(scheduler.config) + if self.config.generation.scheduler == "dpmpp_2m": + options.update(algorithm_type="dpmsolver++", solver_order=2) + return classes[self.config.generation.scheduler].from_config(options) diff --git a/src/mcp_sd15/preprocessors/__init__.py b/src/mcp_sd15/preprocessors/__init__.py new file mode 100644 index 0000000..60e22ba --- /dev/null +++ b/src/mcp_sd15/preprocessors/__init__.py @@ -0,0 +1 @@ +"""Fixed ControlNet preprocessors.""" diff --git a/src/mcp_sd15/preprocessors/depth.py b/src/mcp_sd15/preprocessors/depth.py new file mode 100644 index 0000000..861b889 --- /dev/null +++ b/src/mcp_sd15/preprocessors/depth.py @@ -0,0 +1,24 @@ +"""MiDaS depth-map preprocessing.""" + +from __future__ import annotations + +from typing import cast + +from PIL import Image + + +class DepthPreprocessor: + """Lazily loads the fixed MiDaS annotator from the local cache.""" + + def __init__(self) -> None: + self._detector: object | None = None + + def __call__(self, image: Image.Image) -> Image.Image: + if self._detector is None: + from controlnet_aux import MidasDetector + + self._detector = MidasDetector.from_pretrained( + "lllyasviel/Annotators", local_files_only=True + ) + result = cast(Image.Image, self._detector(image)) # type: ignore[operator] + return result.convert("RGB") diff --git a/src/mcp_sd15/preprocessors/openpose.py b/src/mcp_sd15/preprocessors/openpose.py new file mode 100644 index 0000000..05c879b --- /dev/null +++ b/src/mcp_sd15/preprocessors/openpose.py @@ -0,0 +1,27 @@ +"""OpenPose body or body+hand+face preprocessing.""" + +from __future__ import annotations + +from PIL import Image + + +class OpenPosePreprocessor: + """Lazily loads the fixed OpenPose annotator from the local cache.""" + + def __init__(self) -> None: + self._detector: object | None = None + + def __call__(self, image: Image.Image, *, full: bool) -> Image.Image: + if self._detector is None: + from controlnet_aux import OpenposeDetector + + self._detector = OpenposeDetector.from_pretrained( + "lllyasviel/Annotators", local_files_only=True + ) + result = self._detector( # type: ignore[operator] + image, + include_body=True, + include_hand=full, + include_face=full, + ) + return result.convert("RGB") diff --git a/src/mcp_sd15/schema.py b/src/mcp_sd15/schema.py new file mode 100644 index 0000000..e22450d --- /dev/null +++ b/src/mcp_sd15/schema.py @@ -0,0 +1,142 @@ +"""Validated request and response types independent of MCP and Diffusers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import Any + +from .errors import InvalidRequestError + + +class ImageSize(StrEnum): + SQUARE_512 = "square_512" + LANDSCAPE_512 = "landscape_512" + PORTRAIT_512 = "portrait_512" + + @property + def dimensions(self) -> tuple[int, int]: + return { + ImageSize.SQUARE_512: (512, 512), + ImageSize.LANDSCAPE_512: (512, 384), + ImageSize.PORTRAIT_512: (384, 512), + }[self] + + +class ControlMode(StrEnum): + DEPTH = "depth" + OPENPOSE = "openpose" + OPENPOSE_FULL = "openpose_full" + + +def _prompt(value: Any, label: str, *, required: bool) -> str: + if value is None and not required: + return "" + if not isinstance(value, str): + raise InvalidRequestError(f"{label} must be a string") + value = value.strip() + if required and not value: + raise InvalidRequestError(f"{label} must not be empty") + if len(value) > 10_000: + raise InvalidRequestError(f"{label} is too long") + return value + + +def _size(value: str | ImageSize | None, default: str) -> ImageSize: + try: + return ImageSize(value or default) + except ValueError as exc: + raise InvalidRequestError(f"unsupported size: {value}") from exc + + +@dataclass(frozen=True) +class TextToImageRequest: + positive_prompt: str + negative_prompt: str + size: ImageSize + + @classmethod + def create( + cls, + positive_prompt: str, + negative_prompt: str | None = None, + size: str | ImageSize | None = None, + *, + default_size: str = ImageSize.SQUARE_512, + ) -> TextToImageRequest: + return cls( + _prompt(positive_prompt, "positive_prompt", required=True), + _prompt(negative_prompt, "negative_prompt", required=False), + _size(size, default_size), + ) + + +@dataclass(frozen=True) +class GuidedImageRequest: + positive_prompt: str + negative_prompt: str + size: ImageSize + control_mode: ControlMode + input_image: str + + @classmethod + def create( + cls, + positive_prompt: str, + control_mode: str | ControlMode, + input_image: str, + negative_prompt: str | None = None, + size: str | ImageSize | None = None, + *, + default_size: str = ImageSize.SQUARE_512, + ) -> GuidedImageRequest: + base = TextToImageRequest.create( + positive_prompt, negative_prompt, size, default_size=default_size + ) + try: + mode = ControlMode(control_mode) + except ValueError as exc: + raise InvalidRequestError(f"unsupported control_mode: {control_mode}") from exc + if not isinstance(input_image, str) or not input_image.strip(): + raise InvalidRequestError("input_image must be a non-empty Base64 string") + return cls(base.positive_prompt, base.negative_prompt, base.size, mode, input_image) + + +@dataclass(frozen=True) +class InpaintImageRequest: + positive_prompt: str + negative_prompt: str + size: ImageSize + input_image: str + mask_image: str + + @classmethod + def create( + cls, + positive_prompt: str, + input_image: str, + mask_image: str, + negative_prompt: str | None = None, + size: str | ImageSize | None = None, + *, + default_size: str = ImageSize.SQUARE_512, + ) -> InpaintImageRequest: + base = TextToImageRequest.create( + positive_prompt, negative_prompt, size, default_size=default_size + ) + for label, value in (("input_image", input_image), ("mask_image", mask_image)): + if not isinstance(value, str) or not value.strip(): + raise InvalidRequestError(f"{label} must be a non-empty Base64 string") + return cls( + base.positive_prompt, + base.negative_prompt, + base.size, + input_image, + mask_image, + ) + + +@dataclass(frozen=True) +class GenerationResult: + png_data: bytes + metadata: dict[str, str | int | float] diff --git a/src/mcp_sd15/server.py b/src/mcp_sd15/server.py new file mode 100644 index 0000000..b8dc69d --- /dev/null +++ b/src/mcp_sd15/server.py @@ -0,0 +1,66 @@ +"""MCP public surface: three intentionally small image tools.""" + +from __future__ import annotations + +from typing import Literal + +from mcp.server import MCPServer +from mcp.server.mcpserver import Image + +from .service import ImageGenerationService +from .tools.generate_guided_image import generate_guided_image as guided_adapter +from .tools.generate_image import generate_image as image_adapter +from .tools.inpaint_image import inpaint_image as inpaint_adapter + +Size = Literal["square_512", "landscape_512", "portrait_512"] +GuidanceMode = Literal["depth", "openpose", "openpose_full"] + + +def create_server(service: ImageGenerationService) -> MCPServer: + server = MCPServer("mcp-sd15") + + @server.tool() + async def generate_image( + positive_prompt: str, + negative_prompt: str | None = None, + size: Size | None = None, + ) -> Image: + """Generate one PNG from a positive prompt using the configured SD1.x model.""" + result = await image_adapter(service, positive_prompt, negative_prompt, size) + return Image(data=result.png_data, format="png") + + @server.tool() + async def generate_guided_image( + positive_prompt: str, + control_mode: GuidanceMode, + input_image: str, + negative_prompt: str | None = None, + size: Size | None = None, + ) -> Image: + """Generate one PNG guided by Depth, OpenPose, or full OpenPose. + + input_image is a Base64 PNG/JPEG or an image data URI. + """ + result = await guided_adapter( + service, positive_prompt, control_mode, input_image, negative_prompt, size + ) + return Image(data=result.png_data, format="png") + + @server.tool() + async def inpaint_image( + positive_prompt: str, + input_image: str, + mask_image: str, + negative_prompt: str | None = None, + size: Size | None = None, + ) -> Image: + """Regenerate white mask areas and return one PNG. + + Both images are Base64 PNG/JPEG values or image data URIs. + """ + result = await inpaint_adapter( + service, positive_prompt, input_image, mask_image, negative_prompt, size + ) + return Image(data=result.png_data, format="png") + + return server 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 diff --git a/src/mcp_sd15/tools/__init__.py b/src/mcp_sd15/tools/__init__.py new file mode 100644 index 0000000..0c664f1 --- /dev/null +++ b/src/mcp_sd15/tools/__init__.py @@ -0,0 +1 @@ +"""Thin MCP-independent tool adapters.""" diff --git a/src/mcp_sd15/tools/generate_guided_image.py b/src/mcp_sd15/tools/generate_guided_image.py new file mode 100644 index 0000000..c69286c --- /dev/null +++ b/src/mcp_sd15/tools/generate_guided_image.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from ..schema import GenerationResult, GuidedImageRequest +from ..service import ImageGenerationService + + +async def generate_guided_image( + service: ImageGenerationService, + positive_prompt: str, + control_mode: str, + input_image: str, + negative_prompt: str | None = None, + size: str | None = None, +) -> GenerationResult: + request = GuidedImageRequest.create( + positive_prompt, + control_mode, + input_image, + negative_prompt, + size, + default_size=service.config.generation.default_size, + ) + return await service.generate_guided_image(request) diff --git a/src/mcp_sd15/tools/generate_image.py b/src/mcp_sd15/tools/generate_image.py new file mode 100644 index 0000000..bae30c8 --- /dev/null +++ b/src/mcp_sd15/tools/generate_image.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from ..schema import GenerationResult, TextToImageRequest +from ..service import ImageGenerationService + + +async def generate_image( + service: ImageGenerationService, + positive_prompt: str, + negative_prompt: str | None = None, + size: str | None = None, +) -> GenerationResult: + request = TextToImageRequest.create( + positive_prompt, + negative_prompt, + size, + default_size=service.config.generation.default_size, + ) + return await service.generate_image(request) diff --git a/src/mcp_sd15/tools/inpaint_image.py b/src/mcp_sd15/tools/inpaint_image.py new file mode 100644 index 0000000..be1dbda --- /dev/null +++ b/src/mcp_sd15/tools/inpaint_image.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from ..schema import GenerationResult, InpaintImageRequest +from ..service import ImageGenerationService + + +async def inpaint_image( + service: ImageGenerationService, + positive_prompt: str, + input_image: str, + mask_image: str, + negative_prompt: str | None = None, + size: str | None = None, +) -> GenerationResult: + request = InpaintImageRequest.create( + positive_prompt, + input_image, + mask_image, + negative_prompt, + size, + default_size=service.config.generation.default_size, + ) + return await service.inpaint_image(request) |
