feat: 增加qwen 生图

2026-04-26 11:41:57 +08:00
parent 52479458d2
commit 61b1de2730
7 changed files with 956 additions and 67 deletions
@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+"""
+Qwen Image Generation MCP Server
+Provides image generation via Qwen (DashScope) API
+"""
+
+import os
+import requests
+from urllib.parse import urlparse
+from mcp.server.fastmcp import FastMCP
+
+
+# Initialize MCP server
+mcp = FastMCP("qwen-image-generator")
+
+
+@mcp.tool()
+def generate_image(
+    prompt: str,
+    negative_prompt: str | None = None,
+    prompt_extend: bool = True,
+    size: str = "1024*1024",
+    n: int = 1,
+    image_url: str | None = None,
+    output_path: str | None = None
+) -> str:
+    """
+    Generate images using Qwen (通义千问) image generation API.
+    
+    Args:
+        prompt: The image generation prompt describing what to generate (max 800 chars)
+        negative_prompt: Negative prompt to avoid certain elements (max 500 chars)
+        prompt_extend: Enable prompt extend to enhance the prompt (default: true)
+        watermark: Add watermark to generated image (default: false)
+        size: Image resolution. Options: 1024*1024 (1:1), 1344*768 (16:9), 768*1344 (9:16), 1184*864 (4:3), 864*1184 (3:4)
+        n: Number of images to generate, 1-6 (default: 1)
+        image_url: Optional reference image URL for image-to-image generation
+        output_path: Optional local path to save the generated image
+    
+    Returns:
+        Image URLs and generation status
+    """
+    api_key = os.environ.get("DASHSCOPE_API_KEY", "")
+    api_base = "https://dashscope.aliyuncs.com"
+    
+    if not api_key:
+        return "Error: DASHSCOPE_API_KEY environment variable is not set"
+    
+    # Build content array
+    content = [{"text": prompt}]
+    
+    # Add reference image if provided
+    if image_url:
+        content.append({"image_url": {"url": image_url}})
+    
+    # Build parameters dict
+    parameters = {
+        "prompt_extend": prompt_extend,
+        "size": size,
+        "n": n
+    }
+    
+    # Only add negative_prompt if provided
+    if negative_prompt:
+        parameters["negative_prompt"] = negative_prompt
+    
+    # Build request payload
+    payload = {
+        "model": "qwen-image-2.0-pro",
+        "input": {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": content
+                }
+            ]
+        },
+        "parameters": parameters
+    }
+    
+    # Make API request
+    url = f"{api_base}/api/v1/services/aigc/multimodal-generation/generation"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    
+    try:
+        response = requests.post(url, headers=headers, json=payload, timeout=180)
+        response.raise_for_status()
+        result = response.json()
+        
+        # Check for API errors
+        if "error" in result:
+            error_msg = result.get("error", {}).get("message", "Unknown error")
+            return f"API Error: {error_msg}"
+        
+        # Parse response
+        choices = result.get("output", {}).get("choices", [])
+        usage = result.get("usage", {})
+        
+        # Extract image URLs
+        image_urls = []
+        for choice in choices:
+            message = choice.get("message", {})
+            content_items = message.get("content", [])
+            for item in content_items:
+                if "image" in item:
+                    image_urls.append(item["image"])
+        
+        width = usage.get("width", 1024)
+        height = usage.get("height", 1024)
+        request_id = result.get("request_id", "N/A")
+        
+        # Save image if output_path provided
+        saved_path = None
+        if output_path and image_urls:
+            try:
+                img_response = requests.get(image_urls[0], timeout=30)
+                img_response.raise_for_status()
+                
+                # Determine file extension
+                parsed = urlparse(image_urls[0])
+                ext = os.path.splitext(parsed.path)[1] if "." in parsed.path else ".png"
+                if not ext or len(ext) > 5:
+                    ext = ".png"
+                
+                # Ensure directory exists
+                os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
+                
+                # Add extension if not present
+                if not output_path.endswith(ext):
+                    output_path = output_path + ext
+                
+                with open(output_path, "wb") as f:
+                    f.write(img_response.content)
+                saved_path = os.path.abspath(output_path)
+            except Exception as e:
+                return f"Failed to save image: {str(e)}"
+        
+        # Build response
+        response_text = f"Successfully generated {n} image(s) ({width}x{height})\n\n"
+        response_text += f"Request ID: {request_id}\n\n"
+        
+        if saved_path:
+            response_text += f"Saved to: {saved_path}\n\n"
+        
+        response_text += "Image URLs:\n"
+        for i, img_url in enumerate(image_urls, 1):
+            response_text += f"  {i}. {img_url}\n"
+        
+        return response_text
+        
+    except requests.exceptions.RequestException as e:
+        return f"Request Error: {str(e)}"
+    except Exception as e:
+        return f"Unexpected Error: {str(e)}"
+
+
+if __name__ == "__main__":
+    mcp.run()