diff options
Diffstat (limited to 'src/mcp_sd15')
| -rw-r--r-- | src/mcp_sd15/cli.py | 3 | ||||
| -rw-r--r-- | src/mcp_sd15/model_manager.py | 36 |
2 files changed, 36 insertions, 3 deletions
diff --git a/src/mcp_sd15/cli.py b/src/mcp_sd15/cli.py index abfde01..1248581 100644 --- a/src/mcp_sd15/cli.py +++ b/src/mcp_sd15/cli.py @@ -31,13 +31,14 @@ def main(argv: Sequence[str] | None = None) -> int: try: config = load_config(args.config) logging.getLogger(__name__).info( - "loading checkpoint=%s loras=%d", + "configuration loaded checkpoint=%s loras=%d; models load on generation", config.model.checkpoint_path(), len(config.model.loras), ) backend = ModelManager(config) service = ImageGenerationService(config, backend) server = create_server(service) + logging.getLogger(__name__).info("starting MCP transport=%s", config.server.transport) if config.server.transport == "streamable-http": server.run( transport="streamable-http", diff --git a/src/mcp_sd15/model_manager.py b/src/mcp_sd15/model_manager.py index e87a813..2c314d1 100644 --- a/src/mcp_sd15/model_manager.py +++ b/src/mcp_sd15/model_manager.py @@ -3,6 +3,9 @@ from __future__ import annotations import gc +import logging +import threading +import time from typing import Any import numpy as np @@ -15,21 +18,48 @@ from .preprocessors.depth import DepthPreprocessor from .preprocessors.openpose import OpenPosePreprocessor from .schema import GuidedImageRequest, InpaintImageRequest, TextToImageRequest +logger = logging.getLogger(__name__) + class ModelManager: - """Own all model state; construction validates model compatibility at startup.""" + """Own model state, loading inference models only when generation needs them.""" def __init__(self, config: AppConfig) -> None: self.config = config self.factory = PipelineFactory(config) - self.base = self.factory.create_base() + self.base: Any = None + self._load_lock = threading.Lock() self._control_mode: str | None = None self._controlnet: Any | None = None self._control_pipeline: Any | None = None self._depth = DepthPreprocessor() self._openpose = OpenPosePreprocessor() + def load(self) -> None: + """Load once on success; failed initialization remains retryable. + + Generation calls this from the service's serialized worker thread. + The lock also protects explicit loads by measurement scripts. + """ + with self._load_lock: + if self.base is not None: + return + started = time.monotonic() + logger.info( + "model load started checkpoint=%s loras=%d", + self.config.model.checkpoint_path(), + len(self.config.model.loras), + ) + try: + base = self.factory.create_base() + except Exception: + logger.exception("model load failed") + raise + self.base = base + logger.info("model load completed elapsed_seconds=%.3f", time.monotonic() - started) + def text_to_image(self, request: TextToImageRequest, seed: int) -> Image.Image: + self.load() width, height = request.size.dimensions prompt_embeds, negative_prompt_embeds = self._embed_prompts( request.positive_prompt, request.negative_prompt @@ -49,6 +79,7 @@ class ModelManager: def guided_image( self, request: GuidedImageRequest, image: Image.Image, seed: int ) -> Image.Image: + self.load() width, height = request.size.dimensions image = image.resize((width, height), Image.Resampling.LANCZOS) if request.control_mode == "depth": @@ -79,6 +110,7 @@ class ModelManager: mask: Image.Image, seed: int, ) -> Image.Image: + self.load() width, height = request.size.dimensions image = image.resize((width, height), Image.Resampling.LANCZOS) mask = mask.resize((width, height), Image.Resampling.NEAREST) |