|
|
import os |
|
|
import pathlib |
|
|
import uuid |
|
|
|
|
|
from dotenv import load_dotenv |
|
|
import PIL |
|
|
from google import genai |
|
|
from google.genai import types |
|
|
from mcp.server.fastmcp import FastMCP |
|
|
|
|
|
|
|
|
load_dotenv() |
|
|
|
|
|
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") |
|
|
GEN_IMG_DIR = os.getenv("GEN_IMG_DIR") |
|
|
|
|
|
|
|
|
|
|
|
client = genai.Client(api_key=GEMINI_API_KEY) |
|
|
mcp = FastMCP("virtual_try_on") |
|
|
|
|
|
|
|
|
async def save_image(response, path): |
|
|
for part in response.candidates[0].content.parts: |
|
|
if part.text is not None: |
|
|
continue |
|
|
elif part.inline_data is not None: |
|
|
mime = part.inline_data.mime_type |
|
|
data = part.inline_data.data |
|
|
pathlib.Path(path).write_bytes(data) |
|
|
|
|
|
|
|
|
@mcp.tool() |
|
|
async def try_on(image_path: str, user_prompt: str) -> str: |
|
|
"""Generate a virtual try-on image based on image path and return the saved file path. |
|
|
|
|
|
Args: |
|
|
image_path str: Path to the input image file for try-on image generation |
|
|
user_prompt str: User's request about the try-on image generation. English only |
|
|
|
|
|
Returns: |
|
|
str: File path of the generated image |
|
|
""" |
|
|
try: |
|
|
print(image_path) |
|
|
response = client.models.generate_content( |
|
|
model="models/gemini-2.0-flash-exp", |
|
|
contents=[ |
|
|
"You are a virtual try on tool. Put all fashion items uploaded in the image on a real person and create a picture. " |
|
|
f"{user_prompt}", |
|
|
PIL.Image.open(os.path.abspath(image_path)) |
|
|
], |
|
|
config=types.GenerateContentConfig(response_modalities=['Text', 'Image']) |
|
|
) |
|
|
gen_img_filename = f'{GEN_IMG_DIR}/{uuid.uuid4().hex}.png' |
|
|
await save_image(response, gen_img_filename) |
|
|
return os.path.abspath(gen_img_filename) |
|
|
except Exception as e: |
|
|
print(e) |
|
|
return image_path |
|
|
|
|
|
|
|
|
def main(): |
|
|
print("Started MCP server 'virtual_try_on'...") |
|
|
mcp.run(transport='stdio') |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|
