imgenx-mcp 0.4.0__py3-none-any.whl

imgenx/factory.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+from functools import cache
+from importlib import import_module
+
+from imgenx.predictor.base.base_image_generator import BaseImageGenerator
+from imgenx.predictor.base.base_video_generator import BaseVideoGenerator
+from imgenx.predictor.base.base_image_analyzer import BaseImageAnalyzer
+
+
+@cache
+def create_image_generator(model: str, api_key: str) -> BaseImageGenerator:
+    return create_predictor(model, api_key, 'image_generator')
+
+
+@cache
+def create_video_generator(model: str, api_key: str) -> BaseVideoGenerator:
+    return create_predictor(model, api_key, 'video_generator')
+
+
+@cache
+def create_image_analyzer(model: str, api_key: str) -> BaseImageAnalyzer:
+    return create_predictor(model, api_key, 'image_analyzer')
+
+
+def create_predictor(model: str, api_key: str, role: str) -> BaseImageAnalyzer:
+    provider, model = model.split(':')
+    provider = provider.lower()
+
+    if provider not in get_providers():
+        raise ValueError(f'Provider {provider} not found.')
+
+    predictor_package = f'imgenx.predictor.generators.{provider}_{role.lower()}'
+    predictor_class = f'{provider.capitalize()}{role.split("_")[0].capitalize()}{role.split("_")[1].capitalize()}'
+
+    try:
+        package = import_module(predictor_package)
+        predictor = getattr(package, predictor_class)
+    except (ImportError, AttributeError):
+        raise ValueError(f'Provider {provider} not found.')
+
+    return predictor(model, api_key)
+
+
+@cache
+def get_providers():
+    providers = set()
+    for path in (Path(__file__).parent / 'predictor/generators').glob('*_*_*.py'):
+        providers.add(path.stem.split('_')[0])
+
+    return list(providers)

imgenx/main.py

@@ -0,0 +1,54 @@
+from argparse import ArgumentParser
+
+
+def run():
+    parser = ArgumentParser()
+    subparsers = parser.add_subparsers(dest='command', required=True)
+
+    server_parser = subparsers.add_parser('server', help='启动 MCP 服务器')
+    server_parser.add_argument('--transport', default='stdio', help='stdio|sse|streamable-http')
+    server_parser.add_argument('--host', default='0.0.0.0', help='主机地址')
+    server_parser.add_argument('--port', default=8000, type=int, help='端口')
+    server_parser.add_argument('--no_tools', nargs='+', default=None, help='禁用的工具名列表（用空格分隔）')
+
+    image_parser = subparsers.add_parser('image', help='生成图片')
+    image_parser.add_argument('prompt', help='生成图片的提示词')
+    image_parser.add_argument('--images', nargs='+', default=None, help='输入图片路径列表')
+    image_parser.add_argument('--size', default='2K', help='生成图像的分辨率或宽高像素值，分辨率可选值：1K、2K、4K，宽高像素可选值：2048x2048、2304x1728、1728x2304、2560x1440、1440x2560、2496x1664、1664x2496、3024x1296')
+    image_parser.add_argument('--output', default='imgenx.jpg', help='输出文件或目录路径')
+
+    video_parser = subparsers.add_parser('video', help='生成视频')
+    video_parser.add_argument('prompt', help='生成视频的提示词')
+    video_parser.add_argument('--first_frame', default=None, help='输入视频的第一帧路径')
+    video_parser.add_argument('--last_frame', default=None, help='输入视频的最后一帧路径')
+    video_parser.add_argument('--resolution', default='720p', help='生成视频的分辨率，可选值：480p、720、1080p')
+    video_parser.add_argument('--ratio', default='16:9', help='生成视频的宽高比，可选值：16:9、4:3、1:1、3:4、9:16、21:9')
+    video_parser.add_argument('--duration', default=5, type=int, help='生成视频的时长，单位秒')
+    video_parser.add_argument('--output', default='imgenx.mp4', help='输出文件路径')
+
+    args = parser.parse_args()
+
+    if args.command == 'server':
+        from imgenx.server import mcp
+
+        if args.no_tools:
+            for tool in args.no_tools:
+                mcp.remove_tool(tool)
+
+        if args.transport == 'stdio':
+            mcp.run(transport='stdio')
+        else:
+            mcp.run(transport=args.transport, host=args.host, port=args.port)
+    elif args.command == 'image':
+        from imgenx import script
+        script.gen_image(prompt=args.prompt, size=args.size, output=args.output, images=args.images)
+    elif args.command == 'video':
+        from imgenx import script
+        script.gen_video(prompt=args.prompt, first_frame=args.first_frame, last_frame=args.last_frame,
+                   resolution=args.resolution, ratio=args.ratio, duration=args.duration, output=args.output)
+    else:
+        raise ValueError(f'Unknown command: {args.command}')
+
+
+if __name__ == '__main__':
+    run()

imgenx/operator.py

@@ -0,0 +1,151 @@
+import re
+from typing import Tuple
+from pathlib import Path
+from io import BytesIO
+
+import requests
+from PIL import Image, ImageEnhance
+
+
+def download(url: str, output: str):
+    output = Path(output)
+    if output.exists():
+        raise FileExistsError(f'Path {output} already exists.')
+
+    try:
+        response = requests.get(url, timeout=30)
+        response.raise_for_status()
+    except Exception as e:
+        raise requests.RequestException(f'Error: {e}')
+
+    output.write_bytes(response.content)
+
+
+def get_image_info(image: str) -> str:
+    try:
+        img = _load_image(image)
+    except Exception as e:
+        raise ValueError(f'Error loading image: {e}')
+
+    info = {
+        'format': str(img.format),
+        'size': f"{img.width}x{img.height}",
+        'mode': img.mode,
+        'file_size': _format_file_size(len(img.tobytes()))
+    }
+    return info
+
+
+def crop_image(image: str, box: str, output: str):
+    try:
+        x1, y1, x2, y2 = map(float, box.split(','))
+    except Exception:
+        raise ValueError('box must be "x1, y1, x2, y2" integers')
+
+    img = _load_image(image)
+
+    width = img.width
+    height = img.height
+
+    x1 = width * x1
+    y1 = height * y1
+    x2 = width * x2
+    y2 = height * y2
+    x1 = int(max(0, x1))
+    y1 = int(max(0, y1))
+    x2 = int(max(0, x2))
+    y2 = int(max(0, y2))
+
+    cropped = img.crop((x1, y1, x2, y2))
+    _save_image(cropped, output)
+
+
+def resize_image(image: str, size: str, output: str, keep_aspect: bool = True):
+    m = re.match(r'^(\d+)x(\d+)$', size)
+    if not m:
+        raise ValueError('size must be WIDTHxHEIGHT, e.g., "1024x768"')
+
+    w, h = map(int, m.groups())
+
+    img = _load_image(image)
+    if keep_aspect:
+        resized = img.copy()
+        resized.thumbnail((w, h), Image.Resampling.LANCZOS)
+    else:
+        resized = img.resize((w, h), Image.Resampling.LANCZOS)
+
+    _save_image(resized, output)
+
+
+def convert_image(image: str, format: str, output: str, quality: int = 90):
+    fmt = format.upper()
+    if fmt not in ('PNG', 'JPEG', 'JPG', 'WEBP'):
+        raise ValueError('format must be one of PNG/JPEG/JPG/WEBP')
+
+    img = _load_image(image)
+    _save_image(img, output, format_hint=fmt, quality=quality)
+
+
+def adjust_image(image: str, output: str, brightness: float = 1.0, contrast: float = 1.0, saturation: float = 1.0):
+    img = _load_image(image)
+    img = ImageEnhance.Brightness(img).enhance(brightness)
+    img = ImageEnhance.Contrast(img).enhance(contrast)
+    img = ImageEnhance.Color(img).enhance(saturation)
+
+    _save_image(img, output)
+
+
+def paste_image(front_image: str, backgroud_image: str, position: Tuple[int, int], output: str):
+    front_img = _load_image(front_image)
+    back_img = _load_image(backgroud_image)
+
+    back_img.paste(front_img, tuple(position))
+    _save_image(back_img, output)
+
+
+def _load_image(image: str) -> Image.Image:
+    if re.match(r'^https?://', image):
+        try:
+            resp = requests.get(image, timeout=30)
+            resp.raise_for_status()
+        except Exception as e:
+            raise requests.RequestException(f'Error: {e}')
+        return Image.open(BytesIO(resp.content)).convert('RGBA')
+
+    path = Path(image)
+    if not path.exists():
+        raise FileNotFoundError(f'Image not found: {image}')
+    return Image.open(path)
+
+
+def _save_image(img: Image.Image, output: str, format_hint: str | None = None, quality: int = 90) -> None:
+    out = Path(output)
+    if out.exists():
+        raise FileExistsError(f'Path {out} already exists.')
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    fmt = format_hint or (out.suffix[1:].upper() if out.suffix else 'PNG')
+    
+    if fmt.upper() == 'JPG':
+        fmt = 'JPEG'
+    
+    save_kwargs = {}
+
+    if fmt.upper() in ('JPEG', 'WEBP'):
+        save_kwargs['quality'] = quality
+
+    if fmt.upper() == 'JPEG':
+        img = img.convert('RGB')
+
+    img.save(out, format=fmt, **save_kwargs)
+
+
+def _format_file_size(size_bytes: int) -> str:
+    if size_bytes < 1024:
+        return f"{size_bytes} B"
+    elif size_bytes < 1024 * 1024:
+        return f"{size_bytes / 1024:.1f} KB"
+    elif size_bytes < 1024 * 1024 * 1024:
+        return f"{size_bytes / (1024 * 1024):.1f} MB"
+    else:
+        return f"{size_bytes / (1024 * 1024 * 1024):.1f} GB"

imgenx/oss_service.py

@@ -0,0 +1,198 @@
+"""阿里云 OSS 上传服务"""
+import os
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+import oss2
+from dotenv import load_dotenv
+
+
+load_dotenv()
+
+
+class OSSService:
+    """阿里云 OSS 服务类"""
+
+    def __init__(self):
+        """初始化 OSS 客户端"""
+        self.access_key_id = os.getenv('OSS_ACCESS_KEY_ID')
+        self.access_key_secret = os.getenv('OSS_ACCESS_KEY_SECRET')
+        self.bucket_name = os.getenv('OSS_BUCKET')
+        self.endpoint = os.getenv('OSS_ENDPOINT')
+        self.cdn_url = os.getenv('OSS_CDN_URL', '').rstrip('/')
+
+        if not all([self.access_key_id, self.access_key_secret, self.bucket_name, self.endpoint]):
+            raise ValueError('OSS 配置不完整，请检查环境变量')
+
+        # 创建认证对象
+        auth = oss2.Auth(self.access_key_id, self.access_key_secret)
+
+        # 创建 Bucket 对象
+        self.bucket = oss2.Bucket(auth, self.endpoint, self.bucket_name)
+
+    def generate_object_key(self, filename: str, business_dir: str = 'data') -> str:
+        """
+        生成 OSS 对象键（存储路径）
+
+        格式: /{business_dir}/{YYYYMM}/{timestamp}_{uuid}.{ext}
+
+        Args:
+            filename: 原始文件名
+            business_dir: 业务目录，默认 'data'
+
+        Returns:
+            str: OSS 对象键
+        """
+        # 获取文件扩展名
+        ext = Path(filename).suffix.lower()
+
+        # 生成年月目录
+        now = datetime.now()
+        date_dir = now.strftime('%Y%m')
+
+        # 生成唯一文件名
+        timestamp = int(now.timestamp() * 1000)
+        unique_id = uuid.uuid4().hex[:8]
+        new_filename = f'{timestamp}_{unique_id}{ext}'
+
+        # 组合完整路径
+        object_key = f'{business_dir}/{date_dir}/{new_filename}'
+
+        return object_key
+
+    def upload_file(self, file_path: str, object_key: Optional[str] = None,
+                   business_dir: str = 'data') -> dict:
+        """
+        上传本地文件到 OSS
+
+        Args:
+            file_path: 本地文件路径
+            object_key: OSS 对象键，如果为 None 则自动生成
+            business_dir: 业务目录，默认 'data'
+
+        Returns:
+            dict: 上传结果，包含 object_key, url, cdn_url
+        """
+        file_path = Path(file_path)
+
+        if not file_path.exists():
+            raise FileNotFoundError(f'文件不存在: {file_path}')
+
+        # 如果未指定 object_key，则自动生成
+        if object_key is None:
+            object_key = self.generate_object_key(file_path.name, business_dir)
+
+        # 上传文件
+        with open(file_path, 'rb') as f:
+            result = self.bucket.put_object(object_key, f)
+
+        if result.status != 200:
+            raise Exception(f'上传失败: {result.status}')
+
+        # 生成访问 URL
+        oss_url = f'https://{self.bucket_name}.{self.endpoint}/{object_key}'
+        cdn_url = f'{self.cdn_url}/{object_key}' if self.cdn_url else oss_url
+
+        return {
+            'object_key': object_key,
+            'oss_url': oss_url,
+            'cdn_url': cdn_url,
+            'status': result.status
+        }
+
+    def upload_bytes(self, data: bytes, filename: str,
+                    object_key: Optional[str] = None,
+                    business_dir: str = 'data') -> dict:
+        """
+        上传字节数据到 OSS
+
+        Args:
+            data: 字节数据
+            filename: 文件名（用于提取扩展名）
+            object_key: OSS 对象键，如果为 None 则自动生成
+            business_dir: 业务目录，默认 'data'
+
+        Returns:
+            dict: 上传结果，包含 object_key, url, cdn_url
+        """
+        # 如果未指定 object_key，则自动生成
+        if object_key is None:
+            object_key = self.generate_object_key(filename, business_dir)
+
+        # 上传数据
+        result = self.bucket.put_object(object_key, data)
+
+        if result.status != 200:
+            raise Exception(f'上传失败: {result.status}')
+
+        # 生成访问 URL
+        oss_url = f'https://{self.bucket_name}.{self.endpoint}/{object_key}'
+        cdn_url = f'{self.cdn_url}/{object_key}' if self.cdn_url else oss_url
+
+        return {
+            'object_key': object_key,
+            'oss_url': oss_url,
+            'cdn_url': cdn_url,
+            'status': result.status
+        }
+
+    def delete_file(self, object_key: str) -> bool:
+        """
+        删除 OSS 文件
+
+        Args:
+            object_key: OSS 对象键
+
+        Returns:
+            bool: 删除是否成功
+        """
+        try:
+            result = self.bucket.delete_object(object_key)
+            return result.status == 204
+        except Exception as e:
+            print(f'删除文件失败: {e}')
+            return False
+
+    def file_exists(self, object_key: str) -> bool:
+        """
+        检查文件是否存在
+
+        Args:
+            object_key: OSS 对象键
+
+        Returns:
+            bool: 文件是否存在
+        """
+        try:
+            return self.bucket.object_exists(object_key)
+        except Exception:
+            return False
+
+    def get_file_url(self, object_key: str, use_cdn: bool = True) -> str:
+        """
+        获取文件访问 URL
+
+        Args:
+            object_key: OSS 对象键
+            use_cdn: 是否使用 CDN URL
+
+        Returns:
+            str: 文件访问 URL
+        """
+        if use_cdn and self.cdn_url:
+            return f'{self.cdn_url}/{object_key}'
+        return f'https://{self.bucket_name}.{self.endpoint}/{object_key}'
+
+
+# 创建全局单例
+_oss_service_instance = None
+
+
+def get_oss_service() -> OSSService:
+    """获取 OSS 服务单例"""
+    global _oss_service_instance
+    if _oss_service_instance is None:
+        _oss_service_instance = OSSService()
+    return _oss_service_instance

imgenx/predictor/base/base_image_analyzer.py

@@ -0,0 +1,13 @@
+from typing import List, Dict
+from abc import ABC, abstractmethod
+
+
+class BaseImageAnalyzer(ABC):
+
+    @abstractmethod
+    def __init__(self, model: str, api_key: str):
+        raise NotImplementedError
+
+    @abstractmethod
+    def analyze(self, prompt: str, image: str) -> List[Dict[str, str]]:
+        raise NotImplementedError

imgenx/predictor/base/base_image_generator.py

@@ -0,0 +1,17 @@
+from typing import List, Dict
+from abc import ABC, abstractmethod
+
+
+class BaseImageGenerator(ABC):
+
+    @abstractmethod
+    def __init__(self, model: str, api_key: str):
+        raise NotImplementedError
+
+    @abstractmethod
+    def text_to_image(self, prompt: str, size: str) -> List[Dict[str, str]]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def image_to_image(self, prompt: str, images: List[str], size: str) -> List[Dict[str, str]]:
+        raise NotImplementedError

imgenx/predictor/base/base_video_generator.py

@@ -0,0 +1,18 @@
+from typing import List, Dict
+from abc import ABC, abstractmethod
+
+
+class BaseVideoGenerator(ABC):
+
+    @abstractmethod
+    def __init__(self, model: str, api_key: str):
+        raise NotImplementedError
+
+    @abstractmethod
+    def text_to_video(self, prompt: str, resolution: str, ratio: str, duration: int) -> Dict[str, str]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def image_to_video(self, prompt: str, first_frame: str, last_frame: str|None, 
+                       resolution: str, ratio: str, duration: int) -> Dict[str, str]:
+        raise NotImplementedError

imgenx/predictor/generators/doubao_image_analyzer.py

@@ -0,0 +1,65 @@
+import sys
+sys.path.insert(0, '../../..')
+
+import base64
+from pathlib import Path
+from typing import List, Dict
+from volcenginesdkarkruntime import Ark
+
+from imgenx.predictor.base.base_image_analyzer import BaseImageAnalyzer
+
+
+class DoubaoImageAnalyzer(BaseImageAnalyzer):
+
+    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 analyze(self, prompt: str, image: str) -> str:
+        if not image.startswith('http'):
+            image = self._image_to_base64(image)
+
+        response = self.client.chat.completions.create(
+            model=self.model,   
+            messages=[
+                {
+                    'role': 'user',
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": image
+                            },
+                        },
+                        {"type": "text", "text": prompt},
+                    ],
+                }
+            ],
+        )
+        return response.choices[0].message.content
+
+    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-seed-1-6-vision-250815' 
+
+    analyzer = DoubaoImageAnalyzer(model, api_key)
+    result = analyzer.analyze('请描述这张图片', '/Volumes/DATA/个人/project/imgenx-mcp-server/logo.jpg')
+    print(result)

imgenx/predictor/generators/doubao_image_generator.py

@@ -0,0 +1,77 @@
+import base64
+from pathlib import Path
+from typing import List, Dict
+from volcenginesdkarkruntime import Ark
+
+from imgenx.predictor.base.base_image_generator import BaseImageGenerator
+
+
+class DoubaoImageGenerator(BaseImageGenerator):
+
+    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_image(self, prompt: str, size: str) -> List[Dict[str, str]]:
+        response = self.client.images.generate( 
+            model=self.model,
+            prompt=prompt,
+            sequential_image_generation='auto',
+            response_format='url',
+            size=size,
+            stream=False,
+            watermark=False
+        ) 
+
+        result = []
+        for item in response.model_dump()['data']:
+            if 'b64_json' in item:
+                item.pop('b64_json')
+
+            result.append(item)
+
+        return result
+
+    def image_to_image(self, prompt: str, images: List[str], size: str) -> List[Dict[str, str]]:
+        if len(images) == 0:
+            pass_images = None
+        else:
+            pass_images = []
+
+        for i in images:
+            if i.startswith('http'):
+                pass_images.append(i)
+            else:
+                pass_images.append(self._image_to_base64(i))
+
+        response = self.client.images.generate( 
+            model=self.model,
+            prompt=prompt,
+            sequential_image_generation='auto',
+            response_format='url',
+            image=pass_images,
+            size=size,
+            stream=False,
+            watermark=False
+        ) 
+
+        result = []
+        for item in response.model_dump()['data']:
+            if 'b64_json' in item:
+                item.pop('b64_json')
+
+            result.append(item)
+
+        return result
+
+    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