"""
即梦AI图像生成 MCP 服务器
使用 fastmcp 框架提供图像生成服务
"""
import os
import logging
import re
import traceback
from typing import Optional
from datetime import datetime
from fastmcp import FastMCP
from service import create_service
def replace_quotes(text):
quote_count = 0
def repl(_):
nonlocal quote_count
quote_count += 1
return '“' if quote_count % 2 == 1 else '”'
return re.sub(r'"', repl, text)
# 配置日志
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("jimeng-mcp")
# 创建 FastMCP 实例
mcp = FastMCP("jimeng-image-generator")
# 从环境变量获取token
JIMENG_TOKEN = os.getenv("JIMENG_TOKEN", "")
# 启动日志
logger.info("="*60)
logger.info("即梦AI图像生成 MCP 服务器启动")
logger.info(f"Token状态: {'已配置' if JIMENG_TOKEN else '未配置'}")
logger.info("="*60)
@mcp.tool()
def generate_image(
prompt: str,
save_folder: str,
width: int = 1728,
height: int = 2304,
model: Optional[str] = None,
negative_prompt: Optional[str] = None,
sample_strength: Optional[float] = None
) -> dict:
"""
生成AI图像并保存到指定文件夹
Args:
prompt: 图像生成提示词，详细描述要生成的图像内容
save_folder: 图片保存文件夹路径（绝对路径）
width: 图像宽度（像素），默认1728
height: 图像高度（像素），默认2304
model: 使用的模型，默认为 high_aes_general_v40
negative_prompt: 负面提示词，描述不希望出现的内容
sample_strength: 采样强度，范围0-1，默认0.5
Returns:
dict: 包含生成和下载结果的字典
- success: bool, 是否成功
- submit_id: str, 提交ID
- image_urls: list, 图片URL列表
- saved_files: list, 已保存的文件路径列表
- total_images: int, 图片总数
- save_folder: str, 保存文件夹路径
- error: str, 错误信息（失败时）
示例:
>>> generate_image(
... prompt="秋日街道旁的时尚店铺，一位长发微卷的亚洲女性站在橱窗边",
... save_folder="/tmp/output/images",
... )
"""
# 记录请求信息
prompt = replace_quotes(prompt)
logger.info("="*60)
logger.info("收到图像生成请求")
logger.info(f"提示词: {prompt}")
logger.info(f"保存路径: {save_folder}")
logger.info(f"图像尺寸: {width}x{height}")
if model:
logger.info(f"指定模型: {model}")
if negative_prompt:
logger.info(f"负面提示词: {negative_prompt[:30]}...")
logger.info("="*60)
if not JIMENG_TOKEN:
logger.error("Token未配置")
return {
"success": False,
"error": "JIMENG_TOKEN环境变量未设置，请先配置token"
}
try:
# 创建服务实例
logger.info("创建服务实例...")
service = create_service(JIMENG_TOKEN)
# 调用生成图像方法
logger.info("开始生成图像...")
start_time = datetime.now()
result = service.generate_image(
prompt=prompt,
save_folder=save_folder,
model=model,
negative_prompt=negative_prompt,
sample_strength=sample_strength
)
# 计算耗时
elapsed_time = (datetime.now() - start_time).total_seconds()
# 记录结果
if result.get("success"):
logger.info(f"✅ 图像生成成功！耗时: {elapsed_time:.2f}秒")
logger.info(f"提交ID: {result.get('submit_id')}")
logger.info(f"生成图片数: {result.get('total_images')}")
logger.info(f"保存路径: {result.get('save_folder')}")
saved_files = result.get('saved_files', [])
for idx, file_path in enumerate(saved_files, 1):
logger.info(f" [{idx}] {file_path}")
if result.get('failed_downloads'):
logger.warning(f"部分下载失败: {len(result['failed_downloads'])} 个")
else:
logger.error(f"❌ 图像生成失败: {result.get('error')} {traceback.print_exc()}")
logger.error(f"提交ID: {result.get('submit_id')}")
logger.info("="*60)
return result
except Exception as e:
logger.exception(f"服务调用异常: {str(e)}")
logger.info("="*60)
return {
"success": False,
"error": f"服务调用失败: {str(e)}"
}
@mcp.resource("jimeng://config")
def get_config() -> str:
"""
获取当前配置信息
Returns:
str: 配置信息的文本描述
"""
has_token = bool(JIMENG_TOKEN)
token_preview = f"{JIMENG_TOKEN[:10]}..." if has_token else "未设置"
return f"""即梦AI图像生成服务配置
Token状态: {'已配置' if has_token else '未配置'}
Token预览: {token_preview}
默认配置:
- 默认宽度: 1728px
- 默认高度: 2304px
- 默认模型: high_aes_general_v40
- 采样强度: 0.5
- 轮询间隔: 3秒
- 最大轮询时间: 3分钟
使用说明:
1. 确保设置了 JIMENG_TOKEN 环境变量
2. 调用 generate_image 工具生成图像
3. 图像会自动下载到指定的文件夹
4. 文件按照顺序命名为 01.jpg, 02.jpg 等
支持的图像尺寸:
- 竖图: 1728x2304 (15:16) - 推荐用于小红书
- 竖图: 960x1024 (15:16)
- 横图: 自定义宽高比
工作流程:
1. 提交图像生成请求
2. 自动轮询等待生成完成
3. 获取生成的图片URL
4. 下载并保存到指定文件夹
"""
if __name__ == "__main__":
# 运行MCP服务器
mcp.run()
