diff options
Diffstat (limited to 'src/mcp_sd15')
| -rw-r--r-- | src/mcp_sd15/config.py | 14 | ||||
| -rw-r--r-- | src/mcp_sd15/model_manager.py | 34 |
2 files changed, 41 insertions, 7 deletions
diff --git a/src/mcp_sd15/config.py b/src/mcp_sd15/config.py index e41bdfd..39f31c9 100644 --- a/src/mcp_sd15/config.py +++ b/src/mcp_sd15/config.py @@ -44,6 +44,8 @@ class GenerationConfig: cfg_scale: float = 7.0 scheduler: str = "dpmpp_2m_karras" default_size: str = "square_512" + clip_skip: int = 0 + pad_last_block: bool = False @dataclass(frozen=True) @@ -218,6 +220,14 @@ def load_config(path: str | Path, *, validate_paths: bool = True) -> AppConfig: 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}") + clip_skip_value = generation_raw.get("clip_skip", 0) + if ( + isinstance(clip_skip_value, bool) + or not isinstance(clip_skip_value, int) + or not 0 <= clip_skip_value <= 12 + ): + raise ConfigError("generation.clip_skip must be an integer between 0 and 12") + pad_last_block = _bool(generation_raw, "pad_last_block", False) control_raw = _table(raw, "controlnet") controlnet = ControlNetConfig( @@ -258,7 +268,9 @@ def load_config(path: str | Path, *, validate_paths: bool = True) -> AppConfig: config = AppConfig( model=ModelConfig(root, checkpoint, vae, tuple(loras)), - generation=GenerationConfig(steps_value, cfg_scale, scheduler, default_size), + generation=GenerationConfig( + steps_value, cfg_scale, scheduler, default_size, clip_skip_value, pad_last_block + ), controlnet=controlnet, output=OutputConfig(output_dir.resolve(), _bool(output_raw, "enabled", False)), server=ServerConfig(transport, host, port_value), diff --git a/src/mcp_sd15/model_manager.py b/src/mcp_sd15/model_manager.py index c2bd166..e87a813 100644 --- a/src/mcp_sd15/model_manager.py +++ b/src/mcp_sd15/model_manager.py @@ -31,9 +31,12 @@ class ModelManager: def text_to_image(self, request: TextToImageRequest, seed: int) -> Image.Image: width, height = request.size.dimensions + prompt_embeds, negative_prompt_embeds = self._embed_prompts( + request.positive_prompt, request.negative_prompt + ) result = self.base( - prompt=request.positive_prompt, - negative_prompt=request.negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, width=width, height=height, num_inference_steps=self.config.generation.steps, @@ -53,9 +56,12 @@ class ModelManager: 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=request.positive_prompt, - negative_prompt=request.negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, image=control_image, width=width, height=height, @@ -78,9 +84,12 @@ class ModelManager: 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=request.positive_prompt, - negative_prompt=request.negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, image=image, mask_image=mask, control_image=control_image, @@ -140,3 +149,16 @@ class ModelManager: 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 |