245 lines
8.3 KiB
Python
245 lines
8.3 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
FLUX风格塑形MCP客户端演示脚本
|
||
--------------------------
|
||
演示如何使用MCP客户端连接到服务器并使用FLUX风格塑形功能
|
||
"""
|
||
|
||
import asyncio
|
||
import argparse
|
||
import logging
|
||
import sys
|
||
import os
|
||
import base64
|
||
from typing import Dict, Any
|
||
from pathlib import Path
|
||
|
||
# 设置日志
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 尝试导入MCP客户端库
|
||
try:
|
||
from mcp import ClientSession, HttpServerParameters, types
|
||
from mcp.client.http import http_client
|
||
except ImportError:
|
||
logger.error("未找到MCP客户端库。请安装Model Context Protocol (MCP) Python SDK。")
|
||
logger.error("安装命令: pip install mcp")
|
||
sys.exit(1)
|
||
|
||
async def save_base64_image(base64_data: str, output_path: str) -> str:
|
||
"""
|
||
将base64编码的图像数据保存到文件
|
||
|
||
参数:
|
||
base64_data: 包含MIME类型的base64编码图像数据
|
||
output_path: 输出目录路径
|
||
|
||
返回:
|
||
str: 保存的文件路径
|
||
"""
|
||
# 确保输出目录存在
|
||
os.makedirs(output_path, exist_ok=True)
|
||
|
||
# 从base64字符串提取MIME类型和数据
|
||
if ";" in base64_data and "," in base64_data:
|
||
mime_part, data_part = base64_data.split(",", 1)
|
||
content_type = mime_part.split(";")[0].split(":")[1]
|
||
|
||
# 确定文件扩展名
|
||
if content_type == "image/jpeg":
|
||
ext = ".jpg"
|
||
elif content_type == "image/png":
|
||
ext = ".png"
|
||
elif content_type == "image/webp":
|
||
ext = ".webp"
|
||
else:
|
||
ext = ".png" # 默认扩展名
|
||
else:
|
||
# 如果格式不正确，假设为PNG
|
||
data_part = base64_data
|
||
ext = ".png"
|
||
|
||
# 解码base64数据
|
||
image_data = base64.b64decode(data_part)
|
||
|
||
# 创建输出文件名
|
||
timestamp = asyncio.get_event_loop().time()
|
||
filename = f"flux_output_{int(timestamp)}{ext}"
|
||
file_path = os.path.join(output_path, filename)
|
||
|
||
# 保存图像
|
||
with open(file_path, "wb") as f:
|
||
f.write(image_data)
|
||
|
||
logger.info(f"图像已保存到: {file_path}")
|
||
return file_path
|
||
|
||
async def read_image_to_base64(image_path: str) -> str:
|
||
"""
|
||
读取图像文件并转换为base64编码
|
||
|
||
参数:
|
||
image_path: 图像文件路径
|
||
|
||
返回:
|
||
str: base64编码的图像数据
|
||
"""
|
||
if not os.path.exists(image_path):
|
||
raise FileNotFoundError(f"图像文件不存在: {image_path}")
|
||
|
||
# 读取图像文件
|
||
with open(image_path, "rb") as f:
|
||
image_data = f.read()
|
||
|
||
# 确定MIME类型
|
||
ext = os.path.splitext(image_path)[1].lower()
|
||
if ext in [".jpg", ".jpeg"]:
|
||
mime_type = "image/jpeg"
|
||
elif ext == ".png":
|
||
mime_type = "image/png"
|
||
elif ext == ".webp":
|
||
mime_type = "image/webp"
|
||
else:
|
||
mime_type = "image/png" # 默认MIME类型
|
||
|
||
# 编码为base64
|
||
encoded = base64.b64encode(image_data).decode("utf-8")
|
||
return f"data:{mime_type};base64,{encoded}"
|
||
|
||
async def run_demo(host: str, port: int, structure_image: str, style_image: str, prompt: str,
|
||
depth_strength: float, style_strength: float, output_dir: str):
|
||
"""
|
||
运行FLUX风格塑形演示
|
||
|
||
参数:
|
||
host: 服务器主机地址
|
||
port: 服务器端口
|
||
structure_image: 结构图像路径
|
||
style_image: 风格图像路径
|
||
prompt: 文本提示
|
||
depth_strength: 深度强度
|
||
style_strength: 风格强度
|
||
output_dir: 输出目录
|
||
"""
|
||
# 创建服务器参数
|
||
server_params = HttpServerParameters(
|
||
host=host,
|
||
port=port
|
||
)
|
||
|
||
# 连接到服务器
|
||
logger.info(f"正在连接到MCP服务器: {host}:{port}")
|
||
async with http_client(server_params) as (read, write):
|
||
# 创建客户端会话
|
||
async with ClientSession(read, write) as session:
|
||
# 初始化连接
|
||
logger.info("正在初始化MCP客户端会话...")
|
||
await session.initialize()
|
||
|
||
# 检查FLUX状态
|
||
try:
|
||
logger.info("正在检查FLUX服务状态...")
|
||
content, mime_type = await session.read_resource("flux://status")
|
||
logger.info(f"FLUX服务状态: {content}")
|
||
except Exception as e:
|
||
logger.error(f"获取FLUX状态失败: {e}")
|
||
return
|
||
|
||
# 列出可用工具
|
||
logger.info("正在获取可用工具列表...")
|
||
tools = await session.list_tools()
|
||
logger.info(f"可用工具: {[tool.name for tool in tools]}")
|
||
|
||
# 转换输入图像为base64
|
||
logger.info("正在读取输入图像...")
|
||
structure_image_base64 = await read_image_to_base64(structure_image)
|
||
style_image_base64 = await read_image_to_base64(style_image)
|
||
|
||
# 调用生成风格化图像工具
|
||
logger.info("正在调用生成风格化图像工具...")
|
||
try:
|
||
result = await session.call_tool(
|
||
"生成风格化图像",
|
||
arguments={
|
||
"prompt": prompt,
|
||
"structure_image_base64": structure_image_base64,
|
||
"style_image_base64": style_image_base64,
|
||
"depth_strength": depth_strength,
|
||
"style_strength": style_strength
|
||
}
|
||
)
|
||
|
||
# 处理结果
|
||
if result.get("status") == "success":
|
||
logger.info("图像生成成功")
|
||
|
||
# 保存生成的图像
|
||
image_base64 = result.get("image_base64")
|
||
if image_base64:
|
||
await save_base64_image(image_base64, output_dir)
|
||
else:
|
||
logger.warning("结果中不包含图像数据")
|
||
else:
|
||
logger.error(f"图像生成失败: {result.get('message', '未知错误')}")
|
||
except Exception as e:
|
||
logger.error(f"调用生成风格化图像工具失败: {e}")
|
||
|
||
async def main_async():
|
||
"""异步主函数"""
|
||
# 解析命令行参数
|
||
parser = argparse.ArgumentParser(description='FLUX风格塑形MCP客户端演示')
|
||
parser.add_argument('--host', type=str, default='127.0.0.1', help='MCP服务器地址')
|
||
parser.add_argument('--port', type=int, default=8189, help='MCP服务器端口')
|
||
parser.add_argument('--structure', type=str, required=True, help='结构图像路径')
|
||
parser.add_argument('--style', type=str, required=True, help='风格图像路径')
|
||
parser.add_argument('--prompt', type=str, default='', help='文本提示')
|
||
parser.add_argument('--depth-strength', type=float, default=15.0, help='深度强度')
|
||
parser.add_argument('--style-strength', type=float, default=0.5, help='风格强度')
|
||
parser.add_argument('--output', type=str, default='./output', help='输出目录')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 检查输入文件是否存在
|
||
if not os.path.exists(args.structure):
|
||
logger.error(f"结构图像文件不存在: {args.structure}")
|
||
return 1
|
||
|
||
if not os.path.exists(args.style):
|
||
logger.error(f"风格图像文件不存在: {args.style}")
|
||
return 1
|
||
|
||
# 运行演示
|
||
try:
|
||
await run_demo(
|
||
host=args.host,
|
||
port=args.port,
|
||
structure_image=args.structure,
|
||
style_image=args.style,
|
||
prompt=args.prompt,
|
||
depth_strength=args.depth_strength,
|
||
style_strength=args.style_strength,
|
||
output_dir=args.output
|
||
)
|
||
return 0
|
||
except KeyboardInterrupt:
|
||
logger.info("用户中断操作")
|
||
return 0
|
||
except Exception as e:
|
||
logger.error(f"演示运行出错: {e}")
|
||
return 1
|
||
|
||
def main():
|
||
"""主函数"""
|
||
try:
|
||
return asyncio.run(main_async())
|
||
except KeyboardInterrupt:
|
||
logger.info("用户中断操作")
|
||
return 0
|
||
except Exception as e:
|
||
logger.error(f"程序运行出错: {e}")
|
||
return 1
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main()) |