diff options
Diffstat (limited to 'src/mcp_sd15/model_manager.py')
| -rw-r--r-- | src/mcp_sd15/model_manager.py | 34 |
1 files changed, 28 insertions, 6 deletions
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 |