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