"""Diffusers-backed generation and on-demand ControlNet management.""" from __future__ import annotations import gc import logging import threading import time 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 logger = logging.getLogger(__name__) class ModelManager: """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: 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 ) result = self.base( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, 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: self.load() 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) prompt_embeds, negative_prompt_embeds = self._embed_prompts( request.positive_prompt, request.negative_prompt ) result = pipeline( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, 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: self.load() 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) prompt_embeds, negative_prompt_embeds = self._embed_prompts( request.positive_prompt, request.negative_prompt ) result = pipeline( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, 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] def _embed_prompts(self, positive: str, negative: str) -> tuple[Any, Any]: """Generate prompt embeds via sd_embed to bypass the 77-token limit.""" from sd_embed.embedding_funcs import get_weighted_text_embeddings_sd15 prompt_embeds, neg_prompt_embeds = get_weighted_text_embeddings_sd15( self.base, prompt=positive, neg_prompt=negative or "", pad_last_block=self.config.generation.pad_last_block, clip_skip=self.config.generation.clip_skip, ) return prompt_embeds, neg_prompt_embeds