1 files changed, 142 insertions, 0 deletions
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]