imgenx-mcp 0.4.0__py3-none-any.whl

imgenx/predictor/generators/doubao_video_generator.py

@@ -0,0 +1,122 @@
+import sys
+sys.path.insert(0, '../../..')
+import base64
+import time
+from pathlib import Path
+from typing import List, Dict
+from volcenginesdkarkruntime import Ark
+
+from imgenx.predictor.base.base_video_generator import BaseVideoGenerator
+
+
+class DoubaoVideoGenerator(BaseVideoGenerator):
+
+    def __init__(self, model: str, api_key: str):
+        self.model = model
+        self.client = Ark(
+            base_url='https://ark.cn-beijing.volces.com/api/v3',
+            api_key=api_key,
+        )
+
+    def text_to_video(self, prompt: str, resolution: str = '720p', ratio: str = '16:9', duration: int = 5) -> str:
+        create_result = self.client.content_generation.tasks.create(
+            model=self.model,
+            content=[
+                {
+                    'type': 'text',
+                    'text': prompt.strip()
+                }
+            ],
+            extra_body={
+                'resolution': resolution,
+                'ratio': ratio,
+                'duration': duration,
+            }
+        )
+        task_id = create_result.id
+
+        while True:
+            get_result = self.client.content_generation.tasks.get(task_id=task_id)
+            status = get_result.status
+            if status == "succeeded":
+                return get_result.content.video_url
+            elif status == "failed":
+                raise Exception(str(get_result.error))
+            else:
+                time.sleep(3)
+
+    def image_to_video(self, prompt: str, first_frame: str, last_frame: str|None = None, 
+                       resolution: str = '720p', ratio: str = '16:9', duration: int = 5) -> str:
+        if not first_frame.startswith('http'):
+            first_frame = self._image_to_base64(first_frame)
+
+        if last_frame is not None and not last_frame.startswith('http'):
+            last_frame = self._image_to_base64(last_frame)
+
+        content = [
+            {
+                'type': 'text',
+                'text': prompt.strip()
+            },
+            {
+                'type': 'image_url',
+                'image_url': {
+                    'url': first_frame
+                },
+                'role': 'first_frame'
+            }
+        ]
+
+        if last_frame is not None:
+            content.append({
+                'type': 'image_url',
+                'image_url': {
+                    'url': last_frame
+                },
+                'role': 'last_frame'
+            })
+
+        create_result = self.client.content_generation.tasks.create(
+            model=self.model,
+            content=content,
+            extra_body={
+                'resolution': resolution,
+                'ratio': ratio,
+                'duration': duration,
+            }
+        )
+        task_id = create_result.id
+
+        while True:
+            get_result = self.client.content_generation.tasks.get(task_id=task_id)
+            status = get_result.status
+            if status == "succeeded":
+                return get_result.content.video_url
+            elif status == "failed":
+                raise Exception(str(get_result.error))
+            else:
+                time.sleep(3)
+
+    def _image_to_base64(self, image_path: str) -> str:
+        image_path = Path(image_path)
+
+        with open(image_path, 'rb') as image_file:
+            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
+            base64_image = f'data:image/{image_path.suffix.strip(".")};base64,{base64_image}'
+
+        return base64_image
+
+
+if __name__ == '__main__':
+    import os
+    from dotenv import load_dotenv
+
+    load_dotenv()
+
+    api_key = os.getenv('IMGENX_API_KEY')
+    model = 'doubao-seedance-1-0-pro-fast-251015' 
+
+    generator = DoubaoVideoGenerator(model, api_key)
+    image = '/Volumes/DATA/个人/project/imgenx-mcp-server/logo.jpg'
+    result = generator.image_to_video('一个人在运动', resolution='720p', ratio='16:9', first_frame=image, duration=5)
+    print(result)

imgenx/script.py

@@ -0,0 +1,103 @@
+import os
+from typing import List, Dict
+from pathlib import Path
+from datetime import datetime
+
+import requests
+from dotenv import load_dotenv
+
+from imgenx import factory
+
+
+def text_to_image(model: str, api_key: str, prompt: str, size: str) -> List[Dict[str, str]]:
+    generator = factory.create_image_generator(model, api_key)
+    url_list = generator.text_to_image(prompt, size)
+    return url_list
+
+
+def image_to_image(model: str, api_key: str, prompt: str, images: List[str], size: str) -> List[Dict[str, str]]:
+    generator = factory.create_image_generator(model, api_key)
+    url_list = generator.image_to_image(prompt, images, size)
+    return url_list
+
+
+def text_to_video(model: str, api_key: str, prompt: str,
+                  resolution: str = '720p', ratio: str = '16:9', duration: int = 5) -> str:
+    generator = factory.create_video_generator(model, api_key)
+    url = generator.text_to_video(prompt, resolution, ratio, duration)
+    return url
+
+
+def image_to_video(model: str, api_key: str, prompt: str, first_frame: str, last_frame: str|None = None,
+                  resolution: str = '720p', ratio: str = '16:9', duration: int = 5) -> str:
+    generator = factory.create_video_generator(model, api_key)
+    url = generator.image_to_video(prompt, first_frame, last_frame, resolution, ratio, duration)
+    return url
+
+
+def gen_image(prompt: str, size: str, output: str, images: List[str] = None):
+    print('Generate images...')
+
+    load_dotenv()
+    model = os.getenv('IMGENX_IMAGE_MODEL')
+    api_key = os.getenv('IMGENX_API_KEY')
+
+    if model is None:
+        raise ValueError('Envrioment variable IMGENX_IMAGE_MODEL is empty.')
+
+    if api_key is None:
+        raise ValueError('Envrioment variable IMGENX_API_KEY is empty.')
+
+    output = Path(output)
+
+    if output.exists() and output.is_file():
+        raise ValueError(f'Output path {output} already exists.')
+
+    if images is not None and len(images) > 0:
+        url_list = image_to_image(model, api_key, prompt, images, size)
+    else:
+        url_list = text_to_image(model, api_key, prompt, size)
+
+    if output.is_dir():
+        path_list = [f'{output}/{datetime.now().strftime("%Y-%m-%d_%H:%M:%S")}_{i + 1}.png' for i in range(len(url_list))]
+    elif len(url_list) == 1:
+        path_list = [output]
+    else:
+        path_list = [f'{output.parent}/{output.stem}_{i + 1}{output.suffix if output.suffix else ".jpg"}' for i in range(len(url_list))]
+    
+    for url_item, path in zip(url_list, path_list):
+        response = requests.get(url_item['url'])
+        Path(path).write_bytes(response.content)
+        print(f'Save image to {path}')
+
+
+def gen_video(prompt: str, first_frame: str|None = None, last_frame: str|None = None,
+              resolution: str = '720p', ratio: str = '16:9', duration: int = 5, output: str = None):
+    print('Generate video...')
+
+    load_dotenv()
+    model = os.getenv('IMGENX_VIDEO_MODEL')
+    api_key = os.getenv('IMGENX_API_KEY')
+
+    if model is None:
+        raise ValueError('Envrioment variable IMGENX_VIDEO_MODEL is empty.')
+
+    if api_key is None:
+        raise ValueError('Envrioment variable IMGENX_API_KEY is empty.')
+
+    if output is None:
+        output = f'{datetime.now().strftime("%Y-%m-%d_%H:%M:%S")}.mp4'
+    else:
+        output = Path(output).with_name(f'{Path(output).stem}_{datetime.now().strftime("%Y-%m-%d_%H:%M:%S")}{Path(output).suffix}')
+
+    if output.exists() and output.is_file():
+        raise ValueError(f'Output path {output} already exists.')
+
+    if first_frame is None and last_frame is None:
+        url = text_to_video(model, api_key, prompt, resolution, ratio, duration)
+    else:
+        url = image_to_video(model, api_key, prompt, first_frame, last_frame, resolution, ratio, duration)
+
+    response = requests.get(url)
+    Path(output).write_bytes(response.content)
+    print(f'Save video to {output}')

imgenx/server.py

@@ -0,0 +1,378 @@
+import os
+import re
+from pathlib import Path
+from typing import List, Dict, Tuple
+from urllib.parse import urlparse
+
+import requests
+from dotenv import load_dotenv
+from fastmcp import FastMCP
+from fastmcp.exceptions import ToolError
+from fastmcp.server.dependencies import get_http_headers
+
+from imgenx import factory
+from imgenx import operator
+from imgenx.oss_service import get_oss_service
+
+
+load_dotenv()
+
+mcp = FastMCP(
+    name='imgenx-mcp-server',
+    instructions='图片视频生成工具，自动上传到 OSS 并返回永久 CDN URL',
+)
+
+
+def _upload_url_to_oss(url: str, business_dir: str = 'images') -> Dict[str, str]:
+    """内部函数：下载 URL 并上传到 OSS"""
+    try:
+        # 下载文件
+        response = requests.get(url, timeout=60)
+        response.raise_for_status()
+
+        # 从 URL 提取文件扩展名
+        path = urlparse(url).path
+        ext = Path(path).suffix or '.jpg'
+        filename = f'generated{ext}'
+
+        # 上传到 OSS
+        oss_service = get_oss_service()
+        result = oss_service.upload_bytes(response.content, filename, business_dir=business_dir)
+
+        return result
+    except Exception as e:
+        raise ToolError(f'上传到 OSS 失败: {e}')
+
+
+@mcp.tool
+def text_to_image(prompt: str, size: str = '2K') -> List[Dict[str, str]]:
+    '''根据输入的提示词生成图片，自动上传到 OSS 并返回永久 CDN URL。
+    确保用 Markdown 格式输出图片，例如：![title](cdn_url)
+
+    Args:
+        prompt (str): 生成图片的提示词
+        size (str): 生成图像的分辨率或宽高像素值
+                    分辨率可选值：'1K'、'2K', '4K'
+                    宽高像素可选值：2048x2048、2304x1728、1728x2304、2560x1440、1440x2560、2496x1664、1664x2496、3024x1296
+
+    Returns:
+        List[Dict[str, str]]: 包含 cdn_url（CDN地址）、oss_url（OSS地址）、object_key（存储路径）
+    '''
+    headers = get_http_headers(include_all=True)
+    model = headers.get('imgenx_image_model', os.getenv('IMGENX_IMAGE_MODEL'))
+    api_key = headers.get('imgenx_api_key', os.getenv('IMGENX_API_KEY'))
+
+    if model is None:
+        raise ToolError('IMGENX_IMAGE_MODEL is None')
+
+    if api_key is None:
+        raise ToolError('IMGENX_API_KEY is None')
+
+    try:
+        # 生成图片
+        generator = factory.create_image_generator(model, api_key)
+        temp_url_list = generator.text_to_image(prompt, size)
+
+        # 自动上传到 OSS
+        result_list = []
+        for item in temp_url_list:
+            temp_url = item.get('url')
+            oss_result = _upload_url_to_oss(temp_url, business_dir='images')
+            result_list.append(oss_result)
+
+        return result_list
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+
+@mcp.tool
+def image_to_image(prompt: str, images: List[str], size: str = '2K') -> List[Dict[str, str]]:
+    '''根据输入的提示词和图片生成新图片，自动上传到 OSS 并返回永久 CDN URL。
+    确保用 Markdown 格式输出图片，例如：![title](cdn_url)
+
+    Args:
+        prompt (str): 生成图片的提示词
+        images (List[str]): 输入图片 url 列表或文件路径列表
+        size (str): 生成图像的分辨率或宽高像素值
+                    分辨率可选值：'1K'、'2K', '4K'
+                    宽高像素可选值：2048x2048、2304x1728、1728x2304、2560x1440、1440x2560、2496x1664、1664x2496、3024x1296
+
+    Returns:
+        List[Dict[str, str]]: 包含 cdn_url（CDN地址）、oss_url（OSS地址）、object_key（存储路径）
+    '''
+    headers = get_http_headers(include_all=True)
+    model = headers.get('imgenx_image_model', os.getenv('IMGENX_IMAGE_MODEL'))
+    api_key = headers.get('imgenx_api_key', os.getenv('IMGENX_API_KEY'))
+
+    if model is None:
+        raise ToolError('IMGENX_IMAGE_MODEL is None')
+
+    if api_key is None:
+        raise ToolError('IMGENX_API_KEY is None')
+
+    try:
+        # 生成图片
+        generator = factory.create_image_generator(model, api_key)
+        temp_url_list = generator.image_to_image(prompt, images, size)
+
+        # 自动上传到 OSS
+        result_list = []
+        for item in temp_url_list:
+            temp_url = item.get('url')
+            oss_result = _upload_url_to_oss(temp_url, business_dir='images')
+            result_list.append(oss_result)
+
+        return result_list
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+
+@mcp.tool
+def text_to_video(prompt: str, resolution: str = '720p', ratio: str = '16:9', duration: int = 5) -> Dict[str, str]:
+    '''根据输入的提示词生成视频，自动上传到 OSS 并返回永久 CDN URL。
+    确保用 Markdown 格式输出视频，例如：[title](cdn_url)
+
+    Args:
+        prompt (str): 生成视频的提示词
+        resolution (str): 生成视频的分辨率：480p、720p、1080p
+        ratio (str): 生成视频的比例：16:9、4:3、1:1、3:4、9:16、21:9
+        duration (int): 生成视频的时长，单位秒，支持 2~12 秒
+
+    Returns:
+        Dict[str, str]: 包含 cdn_url（CDN地址）、oss_url（OSS地址）、object_key（存储路径）
+    '''
+    headers = get_http_headers(include_all=True)
+    model = headers.get('imgenx_video_model', os.getenv('IMGENX_VIDEO_MODEL'))
+    api_key = headers.get('imgenx_api_key', os.getenv('IMGENX_API_KEY'))
+
+    if model is None:
+        raise ToolError('IMGENX_VIDEO_MODEL is None')
+
+    if api_key is None:
+        raise ToolError('IMGENX_API_KEY is None')
+
+    try:
+        # 生成视频
+        generator = factory.create_video_generator(model, api_key)
+        temp_url = generator.text_to_video(prompt, resolution, ratio, duration)
+
+        # 自动上传到 OSS
+        result = _upload_url_to_oss(temp_url, business_dir='videos')
+
+        return result
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+
+@mcp.tool
+def image_to_video(prompt: str, first_frame: str, last_frame: str|None = None,
+                  resolution: str = '720p', ratio: str = '16:9', duration: int = 5) -> Dict[str, str]:
+    '''根据输入的提示词和视频首尾帧图片生成视频，自动上传到 OSS 并返回永久 CDN URL。
+    确保用 Markdown 格式输出视频，例如：[title](cdn_url)
+
+    Args:
+        prompt (str): 生成视频的提示词
+        first_frame (str): 视频的首帧图片 url 或文件路径
+        last_frame (str|None): 视频的尾图片 url 或文件路径，默认 None
+        resolution (str): 生成视频的分辨率：480p、720p、1080p
+        ratio (str): 生成视频的比例：16:9、4:3、1:1、3:4、9:16、21:9
+        duration (int): 生成视频的时长，单位秒，支持 2~12 秒
+
+    Returns:
+        Dict[str, str]: 包含 cdn_url（CDN地址）、oss_url（OSS地址）、object_key（存储路径）
+    '''
+    headers = get_http_headers(include_all=True)
+    model = headers.get('imgenx_video_model', os.getenv('IMGENX_VIDEO_MODEL'))
+    api_key = headers.get('imgenx_api_key', os.getenv('IMGENX_API_KEY'))
+
+    if model is None:
+        raise ToolError('IMGENX_VIDEO_MODEL is None')
+
+    if api_key is None:
+        raise ToolError('IMGENX_API_KEY is None')
+
+    try:
+        # 生成视频
+        generator = factory.create_video_generator(model, api_key)
+        temp_url = generator.image_to_video(prompt, first_frame, last_frame, resolution, ratio, duration)
+
+        # 自动上传到 OSS
+        result = _upload_url_to_oss(temp_url, business_dir='videos')
+
+        return result
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+
+@mcp.tool
+def analyze_image(prompt: str, image: str) -> str:
+    '''分析图片获取精确的信息，确保用户需要分析、编辑、裁剪图片时先调用此工具。
+    确保尽量用精确数字描述图片信息。
+    输出图片裁剪区域时，确保给出精确**小数比例坐标**，坐标为左上角和右下角：x1(left), y1(upper), x2(right), y2(lower)
+
+    Args:
+        prompt (str): 分析图片的提示词
+        image (str): 图片路径或 URL
+
+    Returns:
+        str: 图片分析结果
+    '''
+    headers = get_http_headers(include_all=True)
+    model = headers.get('imgenx_analyzer_model', os.getenv('IMGENX_ANALYZER_MODEL'))
+    api_key = headers.get('imgenx_api_key', os.getenv('IMGENX_API_KEY'))
+
+    if model is None:
+        raise ToolError('IMGENX_ANALYZER_MODEL is None')
+
+    if api_key is None:
+        raise ToolError('IMGENX_API_KEY is None')
+
+    try:
+        info = operator.get_image_info(image)
+        prompt = f'image info: {info}\n\n{prompt}'
+
+        analyzer = factory.create_image_analyzer(model, api_key)
+        result = analyzer.analyze(prompt, image)
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+    return result
+
+
+@mcp.tool
+def get_image_info(image: str) -> Dict[str, str]:
+    '''获取图片分辨率和类型信息，确保用户需要获取图片分辨率和类型信息时调用此工具。
+
+    Args:
+        image (str): 图片路径或 URL
+
+    Returns:
+        Dict[str,str]: 分辨率和类型信息
+    '''
+    try:
+        info = operator.get_image_info(image)
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+    return info
+
+
+@mcp.tool
+def crop_image(image: str, box: str, output: str) -> Dict[str, str]:
+    '''框裁剪图片，确保用户需要裁剪图片时调用此工具。
+    Args:
+        image (str): 图片路径或 URL
+        box (str): 小数比例坐标，x1(left), y1(upper), x2(right), y2(lower)
+        output (str): 输出文件路径（后缀决定格式）
+
+    Returns:
+        Dict[str,str]: 生成图片的 path
+    '''
+    try:
+        operator.crop_image(image, box, output)
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+    p = Path(output).resolve()
+    return {'title': p.name, 'path': str(p)}
+
+
+@mcp.tool
+def resize_image(image: str, size: str, output: str, keep_aspect: bool = True) -> Dict[str, str]:
+    '''调整图片尺寸，确保用户需要调整图片尺寸时调用此工具。
+
+    Args:
+        image (str): 图片路径或 URL
+        size (str): "WIDTHxHEIGHT"
+        output (str): 输出文件路径
+        keep_aspect (bool): 是否保持比例（True 为等比不超过目标尺寸）
+
+    Returns:
+        Dict[str,str]: 生成图片的 path
+    '''
+    try:
+        operator.resize_image(image, size, output, keep_aspect=keep_aspect)
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+    p = Path(output).resolve()
+    return {'title': p.name, 'path': str(p)}
+
+
+@mcp.tool
+def convert_image(image: str, format: str, output: str, quality: int = 90) -> Dict[str, str]:
+    '''格式转换，确保用户需要转换图片格式时调用此工具。
+
+    Args:
+        image (str): 图片路径或 URL
+        format (str): 目标格式：PNG/JPEG/JPG/WEBP
+        output (str): 输出文件路径
+        quality (int): 压缩质量（针对有损格式）
+
+    Returns:
+        Dict[str,str]: 生成图片的 path
+    '''
+    try:
+        operator.convert_image(image, format, output, quality=quality)
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+    p = Path(output).resolve()
+    return {'title': p.name, 'path': str(p)}
+
+
+@mcp.tool
+def adjust_image(image: str, output: str, brightness: float = 1.0, contrast: float = 1.0, saturation: float = 1.0) -> Dict[str, str]:
+    '''基础图像调整：亮度/对比度/饱和度，确保用户需要调整图片时调用此工具。
+
+    Args:
+        image (str): 图片路径或 URL
+        output (str): 输出文件路径
+        brightness (float): 亮度，默认 1.0
+        contrast (float): 对比度，默认 1.0
+        saturation (float): 饱和度，默认 1.0
+
+    Returns:
+        Dict[str,str]: 生成图片的 path
+    '''
+    try:
+        operator.adjust_image(image, output, brightness=brightness, contrast=contrast, saturation=saturation)
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+    p = Path(output).resolve()
+    return {'title': p.name, 'path': str(p)}
+
+
+@mcp.tool
+def paste_image(front_image: str, background_image: str, output: str, position: Tuple[int, int]) -> Dict[str, str]:
+    '''将图片粘贴到背景图片上，确保用户需要粘贴图片时调用此工具。
+    粘贴图片前，需要用 `resize_image` 工具调整 `front_image` 到适合的尺寸。
+    调用 `analyze_image` 工具确定粘贴位置。
+
+    Args:
+        front_image (str): 图片路径或 URL
+        background_image (str): 背景图片路径或 URL
+        output (str): 输出文件路径
+        position (Tuple[int, int]): 粘贴位置的 (x, y)，背景图片的左上角像素坐标
+
+    Returns:
+        Dict[str,str]: 生成图片的 path
+    '''
+    try:
+        operator.paste_image(front_image, background_image, position, output)
+    except Exception as e:
+        raise ToolError(f'Error: {e}')
+
+    p = Path(output).resolve()
+    return {'title': p.name, 'path': str(p)}
+
+
+@mcp.custom_route('/health', methods=['GET'])
+def health() -> str:
+    return 'success'
+
+
+@mcp.custom_route('/healthy', methods=['GET'])
+def healthy() -> str:
+    return 'success'