flux1-edit/src/server/models.py

"""Data models for FLUX.1 Edit MCP Server"""

from typing import Optional, Dict, Any, List
from dataclasses import dataclass
from enum import Enum


class ToolName(str, Enum):
    """Available MCP tools"""
    FLUX_EDIT_IMAGE = "flux_edit_image"
    FLUX_EDIT_IMAGE_FROM_FILE = "flux_edit_image_from_file"
    VALIDATE_IMAGE = "validate_image"
    MOVE_TEMP_TO_OUTPUT = "move_temp_to_output"


@dataclass
class ToolParameter:
    """Tool parameter definition"""
    name: str
    type: str
    description: str
    required: bool = True
    enum: Optional[List[str]] = None
    default: Optional[Any] = None


@dataclass
class ToolDefinition:
    """MCP tool definition"""
    name: str
    description: str
    parameters: List[ToolParameter]


# Tool definitions for FLUX.1 Edit MCP Server
TOOL_DEFINITIONS = {
    ToolName.FLUX_EDIT_IMAGE: ToolDefinition(
        name="flux_edit_image",
        description="Edit an image using FLUX.1 Kontext model with base64 input",
        parameters=[
            ToolParameter(
                name="input_image_b64",
                type="string", 
                description="Base64 encoded input image to edit (max 20MB)",
                required=True
            ),
            ToolParameter(
                name="prompt", 
                type="string",
                description="Description of how to edit the image",
                required=True
            ),
            ToolParameter(
                name="seed",
                type="integer",
                description="Seed for reproducible results (0 to 4294967295)",
                required=True
            ),
            ToolParameter(
                name="aspect_ratio",
                type="string",
                description="Image aspect ratio",
                required=False,
                enum=["1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21"],
                default="1:1"
            ),
            ToolParameter(
                name="save_to_file",
                type="boolean",
                description="Whether to save edited image to file",
                required=False,
                default=True
            )
        ]
    ),
    
    ToolName.FLUX_EDIT_IMAGE_FROM_FILE: ToolDefinition(
        name="flux_edit_image_from_file",
        description="Edit an image file from input directory using FLUX.1 Kontext",
        parameters=[
            ToolParameter(
                name="input_image_name",
                type="string",
                description="Name of the image file in input directory",
                required=True
            ),
            ToolParameter(
                name="prompt",
                type="string", 
                description="Description of how to edit the image",
                required=True
            ),
            ToolParameter(
                name="seed",
                type="integer",
                description="Seed for reproducible results (0 to 4294967295)",
                required=True
            ),
            ToolParameter(
                name="aspect_ratio",
                type="string",
                description="Image aspect ratio",
                required=False,
                enum=["1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21"],
                default="1:1"
            ),
            ToolParameter(
                name="save_to_file",
                type="boolean",
                description="Whether to save edited image to file",
                required=False,
                default=True
            )
        ]
    ),
    
    ToolName.VALIDATE_IMAGE: ToolDefinition(
        name="validate_image",
        description="Validate an image file for FLUX.1 Kontext compatibility",
        parameters=[
            ToolParameter(
                name="image_path",
                type="string",
                description="Path to the image file to validate",
                required=True
            )
        ]
    ),
    
    ToolName.MOVE_TEMP_TO_OUTPUT: ToolDefinition(
        name="move_temp_to_output",
        description="Move file from temp directory to output directory",
        parameters=[
            ToolParameter(
                name="temp_file_name",
                type="string",
                description="Name of the file in temp directory to move",
                required=True
            ),
            ToolParameter(
                name="output_file_name",
                type="string",
                description="Desired name for output file (optional)",
                required=False
            ),
            ToolParameter(
                name="copy_only",
                type="boolean",
                description="Copy instead of move (keep original in temp)",
                required=False,
                default=False
            )
        ]
    )
}


@dataclass
class EditResult:
    """Result of image edit operation"""
    success: bool
    base_name: str
    input_file_path: Optional[str] = None
    output_file_path: Optional[str] = None
    json_file_path: Optional[str] = None
    execution_time: float = 0.0
    image_size: Optional[tuple] = None
    error_message: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None


@dataclass 
class ValidationResult:
    """Result of image validation"""
    is_valid: bool
    file_path: str
    size_mb: float
    dimensions: Optional[tuple] = None
    error_message: Optional[str] = None
    warnings: Optional[List[str]] = None


@dataclass
class MoveResult:
    """Result of file move operation"""
    success: bool
    source_path: str
    destination_path: Optional[str] = None
    operation: str = "move"  # "move" or "copy"
    error_message: Optional[str] = None