199 lines
6.7 KiB
Python
199 lines
6.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
MiniMax Image Generation MCP Server
|
|
Provides image generation via MiniMax API
|
|
"""
|
|
|
|
import os
|
|
import base64
|
|
import requests
|
|
from urllib.parse import urlparse
|
|
from mcp.server.fastmcp import FastMCP
|
|
|
|
|
|
# Initialize MCP server
|
|
mcp = FastMCP("minimax-image-generator")
|
|
|
|
|
|
def read_local_image(file_path: str) -> str | None:
|
|
"""
|
|
Read local image file and return as base64 encoded string.
|
|
Returns the base64 string or None if failed.
|
|
"""
|
|
try:
|
|
with open(file_path, "rb") as f:
|
|
image_data = f.read()
|
|
return base64.b64encode(image_data).decode("utf-8")
|
|
except Exception as e:
|
|
return None
|
|
|
|
|
|
def build_subject_reference(subject_ref: str, subject_type: str) -> dict | None:
|
|
"""
|
|
Build subject_reference object from URL or local file path.
|
|
|
|
Args:
|
|
subject_ref: URL or local file path
|
|
subject_type: Type of subject (character, product, logo, etc.)
|
|
|
|
Returns:
|
|
dict with subject_reference structure or None if failed
|
|
"""
|
|
# Check if it's a URL or local file
|
|
if subject_ref.startswith(("http://", "https://")):
|
|
# It's a URL
|
|
return {
|
|
"type": subject_type,
|
|
"image_file": subject_ref
|
|
}
|
|
else:
|
|
# It's a local file - convert to base64
|
|
base64_data = read_local_image(subject_ref)
|
|
if base64_data:
|
|
return {
|
|
"type": subject_type,
|
|
"image_file": f"data:image/jpeg;base64,{base64_data}"
|
|
}
|
|
return None
|
|
|
|
|
|
@mcp.tool()
|
|
def generate_image(
|
|
prompt: str,
|
|
model: str = "image-01",
|
|
aspect_ratio: str = "1:1",
|
|
n: int = 1,
|
|
prompt_optimizer: bool = False,
|
|
subject_reference: str | None = None,
|
|
subject_type: str = "character",
|
|
seed: int | None = None,
|
|
output_path: str | None = None
|
|
) -> str:
|
|
"""
|
|
Generate images using MiniMax image generation API.
|
|
|
|
Args:
|
|
prompt: The image generation prompt describing what to generate
|
|
model: Image generation model (default: image-01)
|
|
aspect_ratio: Image aspect ratio (default: 1:1, options: 1:1, 16:9, 9:16, 4:3, 3:4)
|
|
n: Number of images to generate (default: 1, max: 3)
|
|
prompt_optimizer: Enable prompt optimizer (default: false)
|
|
subject_reference: Reference image URL or local file path for image-to-image generation
|
|
subject_type: Subject reference type (default: character, options: character, product, logo, video_subject, other)
|
|
seed: Random seed for reproducible generation
|
|
output_path: Optional local path to save the generated image
|
|
|
|
Returns:
|
|
Image URLs and generation status
|
|
"""
|
|
api_key = os.environ.get("MINIMAX_API_KEY", "")
|
|
api_base = "https://api.minimaxi.com"
|
|
|
|
if not api_key:
|
|
return "Error: MINIMAX_API_KEY environment variable is not set"
|
|
|
|
# Build request payload
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"aspect_ratio": aspect_ratio,
|
|
"response_format": "url",
|
|
"n": n,
|
|
"prompt_optimizer": prompt_optimizer
|
|
}
|
|
|
|
# Add seed if provided
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
|
|
# Add subject_reference for image-to-image generation
|
|
if subject_reference:
|
|
subject_ref = build_subject_reference(subject_reference, subject_type)
|
|
if subject_ref:
|
|
payload["subject_reference"] = [subject_ref]
|
|
|
|
# Make API request
|
|
url = f"{api_base}/v1/image_generation"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, headers=headers, json=payload, timeout=60)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
|
|
# Check for API errors
|
|
if result.get("base_resp", {}).get("status_code") != 0:
|
|
error_msg = result.get("base_resp", {}).get("status_msg", "Unknown error")
|
|
return f"API Error: {error_msg}"
|
|
|
|
# Extract image URLs
|
|
image_urls = result.get("data", {}).get("image_urls", [])
|
|
metadata = result.get("metadata", {})
|
|
request_id = result.get("id", "N/A")
|
|
success_count = metadata.get("success_count", len(image_urls))
|
|
|
|
# Save image if output_path provided
|
|
saved_paths = []
|
|
if output_path and image_urls:
|
|
try:
|
|
import time
|
|
timestamp = int(time.time())
|
|
|
|
for i, img_url in enumerate(image_urls, 1):
|
|
img_response = requests.get(img_url, timeout=30)
|
|
img_response.raise_for_status()
|
|
|
|
# Determine file extension
|
|
parsed = urlparse(img_url)
|
|
ext = os.path.splitext(parsed.path)[1] if "." in parsed.path else ".jpeg"
|
|
if not ext or len(ext) > 5:
|
|
ext = ".jpeg"
|
|
|
|
# Handle multiple images
|
|
if len(image_urls) > 1:
|
|
base_path = output_path.rsplit('.', 1)[0] if '.' in output_path else output_path
|
|
file_ext = output_path.rsplit('.', 1)[1] if '.' in output_path else ext
|
|
img_path = f"{base_path}_{i}_{timestamp}.{file_ext}"
|
|
else:
|
|
if not output_path.endswith(ext):
|
|
img_path = output_path + ext
|
|
else:
|
|
img_path = output_path
|
|
|
|
# Ensure directory exists
|
|
os.makedirs(os.path.dirname(img_path) if os.path.dirname(img_path) else ".", exist_ok=True)
|
|
|
|
with open(img_path, "wb") as f:
|
|
f.write(img_response.content)
|
|
saved_paths.append(os.path.abspath(img_path))
|
|
except Exception as e:
|
|
return f"Failed to save image: {str(e)}"
|
|
|
|
# Build response
|
|
response_text = f"Successfully generated {success_count} image(s)\n\n"
|
|
response_text += f"Request ID: {request_id}\n\n"
|
|
|
|
if saved_paths:
|
|
response_text += "Saved to:\n"
|
|
for path in saved_paths:
|
|
response_text += f" - {path}\n"
|
|
response_text += "\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()
|