162 lines
5.2 KiB
Python
162 lines
5.2 KiB
Python
#!/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()
|