File size: 2,077 Bytes
f8a73ec 4d9eb9e f8a73ec 4d9eb9e f8a73ec 4d9eb9e f8a73ec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
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")
# os.environ["HTTP_PROXY"] = "http://127.0.0.1:10809"
# os.environ["HTTPS_PROXY"] = "http://127.0.0.1:10809"
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()
|