summaryrefslogtreecommitdiff
path: root/src/mcp_sd15/config.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/config.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/config.py')
-rw-r--r--src/mcp_sd15/config.py237
1 files changed, 237 insertions, 0 deletions
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))