606 lines
24 KiB
Python
606 lines
24 KiB
Python
"""MCP Tool Handlers for FLUX.1 Edit MCP Server"""
|
|
|
|
import json
|
|
import logging
|
|
import random
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
from mcp.types import TextContent, ImageContent
|
|
|
|
from ..connector import Config, FluxEditClient, FluxEditRequest
|
|
from ..utils import (
|
|
validate_edit_parameters,
|
|
validate_file_parameters,
|
|
validate_image_path_parameter,
|
|
validate_move_file_parameters,
|
|
validate_image_file,
|
|
save_image,
|
|
encode_image_base64,
|
|
decode_image_base64,
|
|
sanitize_prompt,
|
|
get_image_dimensions,
|
|
convert_image_to_base64,
|
|
get_file_size_mb
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ToolHandlers:
|
|
"""Handler class for FLUX.1 Edit MCP tools"""
|
|
|
|
def __init__(self, config: Config):
|
|
"""Initialize handlers with configuration"""
|
|
self.config = config
|
|
self.current_seed = None # Track current seed for session
|
|
|
|
def _get_or_create_seed(self) -> int:
|
|
"""Get current seed or create new one"""
|
|
if self.current_seed is None:
|
|
self.current_seed = random.randint(0, 999999)
|
|
return self.current_seed
|
|
|
|
def _reset_seed(self):
|
|
"""Reset seed for new session"""
|
|
self.current_seed = None
|
|
|
|
def _save_b64_to_temp_file(self, b64_data: str, filename: str) -> str:
|
|
"""Save base64 data to a temporary file with specified filename
|
|
|
|
Args:
|
|
b64_data: Base64 encoded image data
|
|
filename: Desired filename for the file
|
|
|
|
Returns:
|
|
str: Path to saved file
|
|
"""
|
|
try:
|
|
# Decode base64 data
|
|
image_data = decode_image_base64(b64_data)
|
|
|
|
# Save to local temp directory for processing
|
|
temp_dir = self.config.base_path / 'temp'
|
|
temp_dir.mkdir(exist_ok=True)
|
|
file_path = temp_dir / filename
|
|
|
|
if not save_image(image_data, str(file_path)):
|
|
raise RuntimeError(f"Failed to save image to temp file: {filename}")
|
|
|
|
logger.info(f"Saved temp file: {filename} ({len(image_data) / 1024:.1f} KB)")
|
|
|
|
return str(file_path)
|
|
except Exception as e:
|
|
logger.error(f"Error saving b64 to temp file: {e}")
|
|
raise
|
|
|
|
def _move_temp_to_generated(self, temp_file_path: str, base_name: str, index: int, extension: str = None) -> str:
|
|
"""
|
|
Move file from temp directory to generated_images directory
|
|
|
|
Args:
|
|
temp_file_path: Path to temporary file
|
|
base_name: Base name for the destination file
|
|
index: Index for the file (0 for input, 1+ for output)
|
|
extension: File extension (will detect from temp file if not provided)
|
|
|
|
Returns:
|
|
str: Path to moved file in generated_images directory
|
|
"""
|
|
try:
|
|
# Ensure output directory exists
|
|
self.config.ensure_output_directory()
|
|
|
|
temp_path = Path(temp_file_path)
|
|
|
|
# Verify source file exists
|
|
if not temp_path.exists():
|
|
raise FileNotFoundError(f"Temp file not found: {temp_file_path}")
|
|
|
|
# Detect extension from temp file if not provided
|
|
if extension is None:
|
|
extension = temp_path.suffix[1:] if temp_path.suffix else 'png'
|
|
|
|
# Generate destination filename
|
|
dest_filename = self.config.generate_filename(base_name, index, extension)
|
|
dest_path = self.config.generated_images_path / dest_filename
|
|
|
|
# Copy file (preserve original in temp for potential reuse)
|
|
import shutil
|
|
try:
|
|
shutil.copy2(temp_file_path, dest_path)
|
|
|
|
# Verify copy was successful
|
|
if not dest_path.exists():
|
|
raise RuntimeError(f"File copy verification failed: {dest_path}")
|
|
|
|
# Check file sizes match
|
|
if temp_path.stat().st_size != dest_path.stat().st_size:
|
|
raise RuntimeError(f"File copy size mismatch: {temp_path.stat().st_size} != {dest_path.stat().st_size}")
|
|
|
|
except PermissionError as e:
|
|
raise RuntimeError(f"Permission denied copying file to {dest_path}: {e}")
|
|
except shutil.Error as e:
|
|
raise RuntimeError(f"Copy operation failed: {e}")
|
|
|
|
logger.info(f"Moved temp file to generated_images: {temp_path.name} → {dest_filename}")
|
|
|
|
return str(dest_path)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error moving temp file to generated_images: {e}")
|
|
raise
|
|
|
|
async def handle_flux_edit_image(self, arguments: Dict[str, Any]) -> List[TextContent | ImageContent]:
|
|
"""
|
|
Handle flux_edit_image tool call
|
|
|
|
Args:
|
|
arguments: Tool arguments
|
|
|
|
Returns:
|
|
List of content items
|
|
"""
|
|
try:
|
|
# Validate parameters
|
|
is_valid, error_msg = validate_edit_parameters(arguments)
|
|
if not is_valid:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Parameter validation failed: {error_msg}"
|
|
)]
|
|
|
|
# Extract parameters
|
|
input_image_b64 = arguments['input_image_b64']
|
|
prompt = sanitize_prompt(arguments['prompt'])
|
|
seed = arguments['seed']
|
|
aspect_ratio = arguments.get('aspect_ratio', self.config.default_aspect_ratio)
|
|
save_to_file = arguments.get('save_to_file', True)
|
|
|
|
logger.info(f"Starting FLUX edit with seed {seed}")
|
|
|
|
# Generate base name
|
|
base_name = self.config.generate_base_name(seed)
|
|
|
|
# Save input image to temp and then to generated_images as 000
|
|
temp_image_name = f'temp_input_{random.randint(1000, 9999)}.png'
|
|
temp_image_path = self._save_b64_to_temp_file(input_image_b64, temp_image_name)
|
|
|
|
# Copy to generated_images as input (000)
|
|
input_generated_path = self._move_temp_to_generated(temp_image_path, base_name, 0)
|
|
logger.info(f"Input file saved: {Path(input_generated_path).name}")
|
|
|
|
# Create FLUX edit request
|
|
request = FluxEditRequest(
|
|
input_image_b64=input_image_b64,
|
|
prompt=prompt,
|
|
seed=seed,
|
|
aspect_ratio=aspect_ratio,
|
|
safety_tolerance=self.config.safety_tolerance,
|
|
output_format=self.config.OUTPUT_FORMAT,
|
|
prompt_upsampling=self.config.prompt_upsampling
|
|
)
|
|
|
|
# Process edit using FLUX API
|
|
async with FluxEditClient(self.config) as client:
|
|
response = await client.edit_image(request)
|
|
|
|
if not response.success:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ FLUX edit failed: {response.error_message}"
|
|
)]
|
|
|
|
# Save output image and metadata
|
|
saved_path = None
|
|
json_path = None
|
|
|
|
if save_to_file:
|
|
output_path = self.config.get_output_path(base_name, 1, 'png')
|
|
|
|
if save_image(response.edited_image_data, str(output_path)):
|
|
saved_path = str(output_path)
|
|
|
|
# Save parameters as JSON
|
|
if self.config.save_parameters:
|
|
params_dict = {
|
|
"base_name": base_name,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"model": self.config.MODEL_NAME,
|
|
"prompt": prompt,
|
|
"seed": seed,
|
|
"aspect_ratio": aspect_ratio,
|
|
"safety_tolerance": self.config.safety_tolerance,
|
|
"output_format": self.config.OUTPUT_FORMAT,
|
|
"prompt_upsampling": self.config.prompt_upsampling,
|
|
"input_image_temp": temp_image_name,
|
|
"input_generated_path": input_generated_path,
|
|
"output_size": response.image_size,
|
|
"execution_time": response.execution_time,
|
|
"request_id": response.request_id,
|
|
"metadata": response.metadata
|
|
}
|
|
|
|
json_path = self.config.get_output_path(base_name, 1, 'json')
|
|
with open(json_path, 'w', encoding='utf-8') as f:
|
|
json.dump(params_dict, f, indent=2, ensure_ascii=False)
|
|
logger.info(f"Parameters saved to: {json_path}")
|
|
|
|
# Prepare response
|
|
contents = []
|
|
|
|
# Add text description
|
|
text = f"✅ Image edited successfully with FLUX.1 Kontext!\n"
|
|
text += f"🎲 Seed: {seed}\n"
|
|
text += f"📁 Base name: {base_name}\n"
|
|
if response.image_size:
|
|
text += f"📐 Size: {response.image_size[0]}x{response.image_size[1]}\n"
|
|
text += f"📏 Aspect ratio: {aspect_ratio}\n"
|
|
text += f"⏱️ Processing time: {response.execution_time:.1f}s\n"
|
|
|
|
if saved_path:
|
|
text += f"\n💾 Output: {Path(saved_path).name}"
|
|
text += f"\n📝 Input: {Path(input_generated_path).name}"
|
|
if json_path:
|
|
text += f"\n📋 Parameters: {Path(json_path).name}"
|
|
|
|
contents.append(TextContent(type="text", text=text))
|
|
|
|
# Add image preview
|
|
if response.edited_image_data:
|
|
image_b64 = encode_image_base64(response.edited_image_data)
|
|
contents.append(ImageContent(
|
|
type="image",
|
|
data=image_b64,
|
|
mimeType="image/png"
|
|
))
|
|
|
|
# Reset seed for next session
|
|
self._reset_seed()
|
|
|
|
return contents
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in handle_flux_edit_image: {e}", exc_info=True)
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Unexpected error: {str(e)}"
|
|
)]
|
|
|
|
async def handle_flux_edit_image_from_file(self, arguments: Dict[str, Any]) -> List[TextContent | ImageContent]:
|
|
"""
|
|
Handle flux_edit_image_from_file tool call
|
|
|
|
Args:
|
|
arguments: Tool arguments
|
|
|
|
Returns:
|
|
List of content items
|
|
"""
|
|
try:
|
|
# Validate parameters
|
|
is_valid, error_msg = validate_file_parameters(arguments)
|
|
if not is_valid:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Parameter validation failed: {error_msg}"
|
|
)]
|
|
|
|
# Extract parameters
|
|
input_image_name = arguments['input_image_name']
|
|
prompt = sanitize_prompt(arguments['prompt'])
|
|
seed = arguments['seed']
|
|
aspect_ratio = arguments.get('aspect_ratio', self.config.default_aspect_ratio)
|
|
save_to_file = arguments.get('save_to_file', True)
|
|
|
|
# Check if file exists in input directory
|
|
input_file_path = self.config.input_path / input_image_name
|
|
|
|
if not input_file_path.exists():
|
|
# Enhanced error message with debug info
|
|
error_text = f"❌ File not found in input directory: {input_image_name}\n"
|
|
error_text += f"📁 Looking in: {self.config.input_path}\n"
|
|
error_text += f"🔍 Full path: {input_file_path}\n"
|
|
error_text += f"📂 Input directory exists: {self.config.input_path.exists()}\n"
|
|
|
|
# List available files in input directory
|
|
if self.config.input_path.exists():
|
|
files = [f.name for f in self.config.input_path.iterdir() if f.is_file()]
|
|
if files:
|
|
error_text += f"📋 Available files: {', '.join(files[:10])}"
|
|
if len(files) > 10:
|
|
error_text += f" and {len(files) - 10} more..."
|
|
else:
|
|
error_text += "📋 No files found in input directory"
|
|
else:
|
|
error_text += "⚠️ Input directory does not exist"
|
|
|
|
return [TextContent(type="text", text=error_text)]
|
|
|
|
# Validate the image file
|
|
is_valid, size_mb, validation_error = validate_image_file(
|
|
str(input_file_path),
|
|
self.config.max_image_size_mb
|
|
)
|
|
if not is_valid:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Image validation failed: {validation_error}"
|
|
)]
|
|
|
|
logger.info(f"Starting FLUX edit from file: {input_image_name} ({size_mb:.2f}MB)")
|
|
|
|
# Convert image to base64
|
|
try:
|
|
input_image_b64 = convert_image_to_base64(str(input_file_path))
|
|
except Exception as e:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Failed to convert image to base64: {str(e)}"
|
|
)]
|
|
|
|
# Generate base name
|
|
base_name = self.config.generate_base_name(seed)
|
|
|
|
# Copy original file to generated_images as input (000)
|
|
try:
|
|
with open(input_file_path, 'rb') as f:
|
|
image_data = f.read()
|
|
|
|
input_generated_path = self.config.get_output_path(base_name, 0, 'png')
|
|
if not save_image(image_data, str(input_generated_path)):
|
|
raise RuntimeError("Failed to save input to generated_images")
|
|
|
|
logger.info(f"Input file copied: {Path(input_generated_path).name}")
|
|
|
|
except Exception as e:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Failed to copy input file: {str(e)}"
|
|
)]
|
|
|
|
# Create FLUX edit request
|
|
request = FluxEditRequest(
|
|
input_image_b64=input_image_b64,
|
|
prompt=prompt,
|
|
seed=seed,
|
|
aspect_ratio=aspect_ratio,
|
|
safety_tolerance=self.config.safety_tolerance,
|
|
output_format=self.config.OUTPUT_FORMAT,
|
|
prompt_upsampling=self.config.prompt_upsampling
|
|
)
|
|
|
|
# Process edit using FLUX API
|
|
async with FluxEditClient(self.config) as client:
|
|
response = await client.edit_image(request)
|
|
|
|
if not response.success:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ FLUX edit failed: {response.error_message}"
|
|
)]
|
|
|
|
# Save output image and metadata
|
|
saved_path = None
|
|
json_path = None
|
|
|
|
if save_to_file:
|
|
output_path = self.config.get_output_path(base_name, 1, 'png')
|
|
|
|
if save_image(response.edited_image_data, str(output_path)):
|
|
saved_path = str(output_path)
|
|
|
|
# Save parameters as JSON
|
|
if self.config.save_parameters:
|
|
params_dict = {
|
|
"base_name": base_name,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"model": self.config.MODEL_NAME,
|
|
"prompt": prompt,
|
|
"seed": seed,
|
|
"aspect_ratio": aspect_ratio,
|
|
"safety_tolerance": self.config.safety_tolerance,
|
|
"output_format": self.config.OUTPUT_FORMAT,
|
|
"prompt_upsampling": self.config.prompt_upsampling,
|
|
"input_image_name": input_image_name,
|
|
"input_file_path": str(input_file_path),
|
|
"input_size": get_image_dimensions(str(input_file_path)),
|
|
"input_size_mb": size_mb,
|
|
"output_size": response.image_size,
|
|
"execution_time": response.execution_time,
|
|
"request_id": response.request_id,
|
|
"metadata": response.metadata
|
|
}
|
|
|
|
json_path = self.config.get_output_path(base_name, 1, 'json')
|
|
with open(json_path, 'w', encoding='utf-8') as f:
|
|
json.dump(params_dict, f, indent=2, ensure_ascii=False)
|
|
logger.info(f"Parameters saved to: {json_path}")
|
|
|
|
# Prepare response
|
|
contents = []
|
|
|
|
# Add text description
|
|
text = f"✅ Image edited successfully from file with FLUX.1 Kontext!\n"
|
|
text += f"📝 Input: {input_image_name} ({size_mb:.2f}MB)\n"
|
|
text += f"🎲 Seed: {seed}\n"
|
|
text += f"📁 Base name: {base_name}\n"
|
|
if response.image_size:
|
|
text += f"📐 Size: {response.image_size[0]}x{response.image_size[1]}\n"
|
|
text += f"📏 Aspect ratio: {aspect_ratio}\n"
|
|
text += f"⏱️ Processing time: {response.execution_time:.1f}s\n"
|
|
|
|
if saved_path:
|
|
text += f"\n💾 Output: {Path(saved_path).name}"
|
|
text += f"\n📝 Input copy: {Path(input_generated_path).name}"
|
|
if json_path:
|
|
text += f"\n📋 Parameters: {Path(json_path).name}"
|
|
|
|
contents.append(TextContent(type="text", text=text))
|
|
|
|
# Add image preview
|
|
if response.edited_image_data:
|
|
image_b64 = encode_image_base64(response.edited_image_data)
|
|
contents.append(ImageContent(
|
|
type="image",
|
|
data=image_b64,
|
|
mimeType="image/png"
|
|
))
|
|
|
|
# Reset seed for next session
|
|
self._reset_seed()
|
|
|
|
return contents
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in handle_flux_edit_image_from_file: {e}", exc_info=True)
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ File-based edit error: {str(e)}"
|
|
)]
|
|
|
|
async def handle_validate_image(self, arguments: Dict[str, Any]) -> List[TextContent]:
|
|
"""
|
|
Handle validate_image tool call
|
|
|
|
Args:
|
|
arguments: Tool arguments
|
|
|
|
Returns:
|
|
List of content items
|
|
"""
|
|
try:
|
|
# Validate parameters
|
|
is_valid, error_msg = validate_image_path_parameter(arguments)
|
|
if not is_valid:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Parameter validation failed: {error_msg}"
|
|
)]
|
|
|
|
image_path = arguments['image_path']
|
|
|
|
# Validate image
|
|
is_valid, size_mb, error_msg = validate_image_file(
|
|
image_path,
|
|
self.config.max_image_size_mb
|
|
)
|
|
|
|
# Get additional info if valid
|
|
if is_valid:
|
|
width, height = get_image_dimensions(image_path)
|
|
|
|
text = f"✅ Image validation passed!\n"
|
|
text += f"📁 File: {Path(image_path).name}\n"
|
|
text += f"📐 Dimensions: {width}x{height}\n"
|
|
text += f"💾 Size: {size_mb:.2f}MB\n"
|
|
text += f"🎯 Max allowed: {self.config.max_image_size_mb}MB\n"
|
|
|
|
# Check aspect ratio compatibility
|
|
from ..utils import get_optimal_aspect_ratio
|
|
optimal_ratio = get_optimal_aspect_ratio(width, height)
|
|
text += f"📏 Optimal aspect ratio: {optimal_ratio}"
|
|
else:
|
|
text = f"❌ Image validation failed!\n"
|
|
text += f"📁 File: {Path(image_path).name}\n"
|
|
text += f"⚠️ Issue: {error_msg}"
|
|
|
|
return [TextContent(type="text", text=text)]
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in handle_validate_image: {e}", exc_info=True)
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Validation error: {str(e)}"
|
|
)]
|
|
|
|
async def handle_move_temp_to_output(self, arguments: Dict[str, Any]) -> List[TextContent]:
|
|
"""
|
|
Handle move_temp_to_output tool call
|
|
|
|
Args:
|
|
arguments: Tool arguments
|
|
|
|
Returns:
|
|
List of content items
|
|
"""
|
|
try:
|
|
# Validate parameters
|
|
is_valid, error_msg = validate_move_file_parameters(arguments)
|
|
if not is_valid:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Parameter validation failed: {error_msg}"
|
|
)]
|
|
|
|
temp_file_name = arguments['temp_file_name']
|
|
output_file_name = arguments.get('output_file_name')
|
|
copy_only = arguments.get('copy_only', False)
|
|
|
|
# Get temp file path
|
|
temp_file_path = self.config.base_path / 'temp' / temp_file_name
|
|
|
|
# Check if temp file exists
|
|
if not temp_file_path.exists():
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Temp file not found: {temp_file_name}"
|
|
)]
|
|
|
|
# Generate output file name if not provided
|
|
if not output_file_name:
|
|
base_name = self.config.generate_base_name_simple()
|
|
file_ext = Path(temp_file_name).suffix[1:] or 'png'
|
|
output_file_name = f"{base_name}_001.{file_ext}"
|
|
|
|
# Ensure output directory exists
|
|
self.config.ensure_output_directory()
|
|
|
|
# Get output path
|
|
output_path = self.config.generated_images_path / output_file_name
|
|
|
|
# Move or copy file
|
|
try:
|
|
import shutil
|
|
if copy_only:
|
|
shutil.copy2(temp_file_path, output_path)
|
|
operation = "copied"
|
|
else:
|
|
shutil.move(str(temp_file_path), str(output_path))
|
|
operation = "moved"
|
|
|
|
# Verify operation was successful
|
|
if not output_path.exists():
|
|
raise RuntimeError(f"File {operation} verification failed")
|
|
|
|
logger.info(f"📁 File {operation}: {temp_file_name} -> {output_file_name}")
|
|
|
|
# Get file size for reporting
|
|
file_size_mb = output_path.stat().st_size / (1024 * 1024)
|
|
|
|
text = f"✅ File {operation} successfully!\n"
|
|
text += f"📁 From temp: {temp_file_name}\n"
|
|
text += f"📁 To output: {output_file_name}\n"
|
|
text += f"💾 Size: {file_size_mb:.2f}MB"
|
|
|
|
return [TextContent(type="text", text=text)]
|
|
|
|
except PermissionError as e:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ Permission denied: {str(e)}"
|
|
)]
|
|
except Exception as e:
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ File operation failed: {str(e)}"
|
|
)]
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in handle_move_temp_to_output: {e}", exc_info=True)
|
|
return [TextContent(
|
|
type="text",
|
|
text=f"❌ File move error: {str(e)}"
|
|
)]
|