"""Midjourney MCP Server - Main entry point with tool functions."""
import logging
from typing import List, Optional
from mcp.server.fastmcp import FastMCP
from service import get_service
from config import setup_logging, get_config
# Initialize logging
setup_logging()
logger = logging.getLogger(__name__)
# Initialize MCP server
mcp = FastMCP("midjourney")
# Get configuration
config = get_config()
# ============================================================================
# MCP Tool Functions - Complete Midjourney Functionality
# ============================================================================
@mcp.tool()
async def imagine_image(
prompt: str,
aspect_ratio: str = "1:1",
base64_images: Optional[List[str]] = None
) -> str:
"""Generate images from text prompts with optional reference images.
Args:
prompt: Text description of the image to generate (English only)
aspect_ratio: Aspect ratio of the image (e.g., "16:9", "1:1", "9:16")
base64_images: Optional list of reference images in base64 format
Returns:
Generated image URL and task information
"""
try:
logger.info(f"Imagine request: {prompt[:100]}...")
service = await get_service()
result = await service.imagine(prompt, aspect_ratio, base64_images)
return result
except Exception as e:
logger.error(f"Error in imagine_image: {e}")
return f"❌ Error generating image: {str(e)}"
@mcp.tool()
async def blend_images(
base64_images: List[str],
dimensions: str = "SQUARE"
) -> str:
"""Blend multiple images together.
Args:
base64_images: List of 2-5 images to blend in base64 format
dimensions: Output dimensions ("PORTRAIT", "SQUARE", "LANDSCAPE")
Returns:
Blended image URL and task information
"""
try:
logger.info(f"Blend request: {len(base64_images)} images")
service = await get_service()
result = await service.blend(base64_images, dimensions)
return result
except Exception as e:
logger.error(f"Error in blend_images: {e}")
return f"❌ Error blending images: {str(e)}"
@mcp.tool()
async def describe_image(base64_image: str) -> str:
"""Generate text descriptions of an image.
Args:
base64_image: Image to describe in base64 format
Returns:
Text description of the image
"""
try:
logger.info("Describe image request")
service = await get_service()
result = await service.describe(base64_image)
return result
except Exception as e:
logger.error(f"Error in describe_image: {e}")
return f"❌ Error describing image: {str(e)}"
@mcp.tool()
async def change_image(
task_id: str,
action: str,
index: Optional[int] = None
) -> str:
"""Create variations, upscales, or rerolls of existing images.
Args:
task_id: ID of the original generation task
action: Action type ("UPSCALE", "VARIATION", "REROLL")
index: Image index (1-4) for UPSCALE and VARIATION actions
Returns:
Modified image URL and task information
"""
try:
logger.info(f"Change request: {action} for task {task_id}")
service = await get_service()
result = await service.change(task_id, action, index)
return result
except Exception as e:
logger.error(f"Error in change_image: {e}")
return f"❌ Error changing image: {str(e)}"
@mcp.tool()
async def modal_edit(
task_id: str,
action: str,
prompt: Optional[str] = None
) -> str:
"""Perform advanced editing like zoom, pan, or inpainting.
Args:
task_id: ID of the original generation task
action: Edit action type (zoom, pan, inpaint, etc.)
prompt: Additional prompt for the edit
Returns:
Edited image URL and task information
"""
try:
logger.info(f"Modal edit request: {action} for task {task_id}")
service = await get_service()
result = await service.modal_edit(task_id, action, prompt)
return result
except Exception as e:
logger.error(f"Error in modal_edit: {e}")
return f"❌ Error in modal edit: {str(e)}"
@mcp.tool()
async def swap_face(source_image: str, target_image: str) -> str:
"""Swap faces between two images.
Args:
source_image: Source face image in base64 format
target_image: Target image in base64 format
Returns:
Face-swapped image URL and task information
"""
try:
logger.info("Face swap request")
service = await get_service()
result = await service.swap_face(source_image, target_image)
return result
except Exception as e:
logger.error(f"Error in swap_face: {e}")
return f"❌ Error swapping faces: {str(e)}"
# ============================================================================
# Task Management Tools
# ============================================================================
@mcp.tool()
async def get_task_status(task_id: str) -> str:
"""Get current status of a Midjourney task.
Args:
task_id: Task ID to check
Returns:
Current task status and details
"""
try:
logger.info(f"Task status request: {task_id}")
service = await get_service()
result = await service.get_task_status(task_id)
return result
except Exception as e:
logger.error(f"Error in get_task_status: {e}")
return f"❌ Error getting task status: {str(e)}"
# ============================================================================
# Server Lifecycle Management
# ============================================================================
# Server lifecycle will be handled by the main function
def main():
"""Main entry point for the MCP server."""
try:
logger.info("Starting Midjourney MCP Server...")
logger.info(f"Configuration: API Key configured = {bool(config.gptnb_api_key)}")
logger.info(f"Base URL: {config.gptnb_base_url}")
logger.info("Server started successfully!")
mcp.run(transport="stdio")
except KeyboardInterrupt:
logger.info("Server interrupted by user")
except Exception as e:
logger.error(f"Server error: {e}")
raise
if __name__ == "__main__":
main()
