diff options
Diffstat (limited to 'src/mcp_sd15/server.py')
| -rw-r--r-- | src/mcp_sd15/server.py | 66 |
1 files changed, 66 insertions, 0 deletions
diff --git a/src/mcp_sd15/server.py b/src/mcp_sd15/server.py new file mode 100644 index 0000000..b8dc69d --- /dev/null +++ b/src/mcp_sd15/server.py @@ -0,0 +1,66 @@ +"""MCP public surface: three intentionally small image tools.""" + +from __future__ import annotations + +from typing import Literal + +from mcp.server import MCPServer +from mcp.server.mcpserver import Image + +from .service import ImageGenerationService +from .tools.generate_guided_image import generate_guided_image as guided_adapter +from .tools.generate_image import generate_image as image_adapter +from .tools.inpaint_image import inpaint_image as inpaint_adapter + +Size = Literal["square_512", "landscape_512", "portrait_512"] +GuidanceMode = Literal["depth", "openpose", "openpose_full"] + + +def create_server(service: ImageGenerationService) -> MCPServer: + server = MCPServer("mcp-sd15") + + @server.tool() + async def generate_image( + positive_prompt: str, + negative_prompt: str | None = None, + size: Size | None = None, + ) -> Image: + """Generate one PNG from a positive prompt using the configured SD1.x model.""" + result = await image_adapter(service, positive_prompt, negative_prompt, size) + return Image(data=result.png_data, format="png") + + @server.tool() + async def generate_guided_image( + positive_prompt: str, + control_mode: GuidanceMode, + input_image: str, + negative_prompt: str | None = None, + size: Size | None = None, + ) -> Image: + """Generate one PNG guided by Depth, OpenPose, or full OpenPose. + + input_image is a Base64 PNG/JPEG or an image data URI. + """ + result = await guided_adapter( + service, positive_prompt, control_mode, input_image, negative_prompt, size + ) + return Image(data=result.png_data, format="png") + + @server.tool() + async def inpaint_image( + positive_prompt: str, + input_image: str, + mask_image: str, + negative_prompt: str | None = None, + size: Size | None = None, + ) -> Image: + """Regenerate white mask areas and return one PNG. + + Both images are Base64 PNG/JPEG values or image data URIs. + """ + result = await inpaint_adapter( + service, positive_prompt, input_image, mask_image, negative_prompt, size + ) + return Image(data=result.png_data, format="png") + + return server |
