From 431d012e20a5cafa864abc3762608b99cf0f5ef4 Mon Sep 17 00:00:00 2001 From: ened Date: Tue, 26 Aug 2025 04:10:31 +0900 Subject: [PATCH] fix bug --- CLAUDE.md | 6 +-- main.py | 26 ++++++++-- src/connector/config.py | 11 ++-- src/connector/flux_client.py | 78 +++++++++++++++++----------- src/server/handlers.py | 99 +----------------------------------- src/server/models.py | 51 ------------------- src/utils/__init__.py | 2 - src/utils/validation.py | 66 +----------------------- 8 files changed, 84 insertions(+), 255 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index e622efd..0bddda9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -69,15 +69,14 @@ python test_fixes.py # 반드시 실행하여 확인 "seed": "필수 입력 (일관된 편집 결과를 위해)", "safety_tolerance": 2, "output_format": "png", - "webhook_url": null, - "webhook_secret": null + "prompt_upsampling": false } ``` #### 2.2 설정 방향 - **prompt**: 토큰 크기 제한이 명시되지 않음 - 제한없이 처리 - **input_image**: 20MB 크기 제한 반영 -- **aspect_ratio**: 기본값 1:1 또는 16:9 설정 (확인 필요) +- **aspect_ratio**: 편집 작업에서는 사용하지 않음 (입력 이미지의 원본 비율 유지) - **seed**: 반드시 입력받아서 일관된 스타일 유지 및 재현성 보장 - **prompt_upsampling**: false (기본) - **safety_tolerance**: 2 (기본값) @@ -135,7 +134,6 @@ D:\Project\little-fairy\flux1-edit\ "input_image_b64": str, # Base64 인코딩된 입력 이미지 "prompt": str, # 편집 설명 "seed": int, # 재현성을 위한 시드값 - "aspect_ratio": str, # "1:1" | "16:9" 등 "save_to_file": bool # 파일 저장 여부 (기본: True) } ``` diff --git a/main.py b/main.py index 2e6d0ec..e6ac8a7 100644 --- a/main.py +++ b/main.py @@ -3,7 +3,7 @@ """ FLUX.1 Edit MCP Server - Fixed Version -FLUX.1 Kontext를 사용한 AI 이미지 편집 MCP 서버 +AI image editing MCP server using FLUX.1 Kontext - Enhanced error handling and UTF-8 support - MCP protocol compliance - Based on imagen4 server structure @@ -256,6 +256,28 @@ class FluxEditMCPServer: logger.error(f"Error listing tools: {e}", exc_info=True) raise + @self.server.list_prompts() + async def handle_list_prompts() -> List[types.Prompt]: + """List available prompts (empty for this server)""" + try: + logger.info("Listing available prompts") + # This server doesn't provide prompt templates + return [] + except Exception as e: + logger.error(f"Error listing prompts: {e}", exc_info=True) + raise + + @self.server.list_resources() + async def handle_list_resources() -> List[types.Resource]: + """List available resources (empty for this server)""" + try: + logger.info("Listing available resources") + # This server doesn't provide static resources + return [] + except Exception as e: + logger.error(f"Error listing resources: {e}", exc_info=True) + raise + @self.server.call_tool() async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[types.TextContent | types.ImageContent]: """Handle tool calls with comprehensive error handling""" @@ -271,8 +293,6 @@ class FluxEditMCPServer: return await self.handlers.handle_flux_edit_image_from_file(arguments) elif name == ToolName.VALIDATE_IMAGE: return await self.handlers.handle_validate_image(arguments) - elif name == ToolName.MOVE_TEMP_TO_OUTPUT: - return await self.handlers.handle_move_temp_to_output(arguments) else: error_msg = f"Unknown tool: {name}" logger.error(error_msg) diff --git a/src/connector/config.py b/src/connector/config.py index 1ee400a..89b30ec 100644 --- a/src/connector/config.py +++ b/src/connector/config.py @@ -16,21 +16,22 @@ class Config: # FLUX.1 Kontext API Configuration API_BASE_URL = "https://api.bfl.ai" - EDIT_ENDPOINT = "/flux-kontext-pro" + EDIT_ENDPOINT = "/v1/flux-kontext-pro" RESULT_ENDPOINT = "/v1/get_result" - # Fixed FLUX parameters based on requirements + # FLUX.1 Kontext default parameters based on official documentation MODEL_NAME = "flux-kontext-pro" - OUTPUT_FORMAT = "png" + OUTPUT_FORMAT = "png" # Fixed to PNG format PROMPT_UPSAMPLING = False DEFAULT_SAFETY_TOLERANCE = 2 # Image size limits MAX_IMAGE_SIZE_MB = 20 # FLUX.1 Kontext limit - # Aspect ratios supported + # Aspect ratios supported (from 3:7 to 7:3 according to docs) SUPPORTED_ASPECT_RATIOS = [ - "1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21" + "1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21", + "3:7", "7:3", "5:7", "7:5" # Extended range as per docs ] DEFAULT_ASPECT_RATIO = "1:1" diff --git a/src/connector/flux_client.py b/src/connector/flux_client.py index 8504ce6..d07e181 100644 --- a/src/connector/flux_client.py +++ b/src/connector/flux_client.py @@ -20,7 +20,6 @@ class FluxEditRequest: input_image_b64: str prompt: str seed: int - aspect_ratio: str = "1:1" safety_tolerance: int = 2 output_format: str = "png" prompt_upsampling: bool = False @@ -59,7 +58,12 @@ class FluxEditClient: async def _ensure_session(self): """Ensure aiohttp session is created""" if self.session is None or self.session.closed: - timeout = aiohttp.ClientTimeout(total=self.config.default_timeout) + # Increase timeout for large image uploads and processing + timeout = aiohttp.ClientTimeout( + total=self.config.default_timeout, + connect=30, # Connection timeout + sock_read=60 # Socket read timeout + ) self.session = aiohttp.ClientSession(timeout=timeout) async def close(self): @@ -71,8 +75,9 @@ class FluxEditClient: def _get_headers(self) -> Dict[str, str]: """Get request headers with API key""" return { - 'Content-Type': 'application/json', - 'X-Key': self.config.api_key + 'accept': 'application/json', + 'x-key': self.config.api_key, + 'Content-Type': 'application/json' } async def _create_edit_request(self, request: FluxEditRequest) -> Optional[str]: @@ -88,54 +93,70 @@ class FluxEditClient: try: await self._ensure_session() - # Prepare request payload + # Prepare request payload for FLUX.1 Kontext API (complete parameters) payload = { "prompt": request.prompt, "input_image": request.input_image_b64, - "seed": request.seed, + "seed": request.seed if request.seed is not None else None, "safety_tolerance": request.safety_tolerance, - "output_format": request.output_format + "output_format": request.output_format, + "prompt_upsampling": request.prompt_upsampling } - # Add optional parameters based on API spec - if hasattr(request, 'aspect_ratio') and request.aspect_ratio: - # Note: Check if FLUX.1 Kontext actually supports aspect_ratio parameter - # This might need to be removed based on actual API response - payload["aspect_ratio"] = request.aspect_ratio + # aspect_ratio is not used in editing operations (preserve input image aspect ratio) + url = self.config.get_api_url(self.config.EDIT_ENDPOINT) logger.info(f"Creating FLUX edit request to {url}") logger.debug(f"Request payload keys: {list(payload.keys())}") + # Log payload size for debugging (without exposing image data) + payload_size = len(str(payload)) + image_size = len(request.input_image_b64) if request.input_image_b64 else 0 + logger.debug(f"Payload size: {payload_size} chars, Image size: {image_size} chars") + async with self.session.post(url, json=payload, headers=self._get_headers()) as response: if response.status == 200: result = await response.json() request_id = result.get('id') - if request_id: + polling_url = result.get('polling_url') + + if request_id and polling_url: logger.info(f"Edit request created successfully: {request_id}") - return request_id + logger.debug(f"Polling URL: {polling_url}") + # Store polling_url for use in polling + return (request_id, polling_url) else: - logger.error(f"No request_id in response: {result}") + logger.error(f"Missing id or polling_url in response: {result}") return None else: - error_text = await response.text() - logger.error(f"Failed to create edit request: {response.status} - {error_text}") + try: + error_text = await response.text() + logger.error(f"Failed to create edit request: {response.status} - {error_text}") + except Exception as text_error: + logger.error(f"Failed to create edit request: {response.status} - Could not read error response: {text_error}") return None except asyncio.TimeoutError: logger.error("Timeout creating edit request") return None + except aiohttp.ClientPayloadError as e: + logger.error(f"Client payload error creating edit request: {e}") + return None + except aiohttp.ClientError as e: + logger.error(f"Client error creating edit request: {e}") + return None except Exception as e: logger.error(f"Error creating edit request: {e}", exc_info=True) return None - async def _poll_result(self, request_id: str) -> Optional[Dict[str, Any]]: + async def _poll_result(self, polling_url: str) -> Optional[Dict[str, Any]]: """ - Poll for edit result using request_id + Poll for edit result using polling_url Args: - request_id: Request ID from create_edit_request + polling_url: Polling URL from create_edit_request Returns: dict: Result data if successful, None otherwise @@ -143,18 +164,16 @@ class FluxEditClient: try: await self._ensure_session() - url = self.config.get_api_url(self.config.RESULT_ENDPOINT) - params = {"id": request_id} - + # Use the provided polling URL directly attempts = 0 max_attempts = self.config.max_polling_attempts interval = self.config.polling_interval - logger.info(f"Starting to poll for result: {request_id}") + logger.info(f"Starting to poll using URL: {polling_url}") while attempts < max_attempts: try: - async with self.session.get(url, params=params, headers=self._get_headers()) as response: + async with self.session.get(polling_url, headers=self._get_headers()) as response: if response.status == 200: result = await response.json() @@ -260,16 +279,18 @@ class FluxEditClient: logger.info(f"Starting FLUX image edit with seed {request.seed}") # Step 1: Create edit request - request_id = await self._create_edit_request(request) - if not request_id: + create_result = await self._create_edit_request(request) + if not create_result: return FluxEditResponse( success=False, error_message="Failed to create edit request", execution_time=(datetime.now() - start_time).total_seconds() ) + request_id, polling_url = create_result + # Step 2: Poll for result - result = await self._poll_result(request_id) + result = await self._poll_result(polling_url) if not result: return FluxEditResponse( success=False, @@ -314,7 +335,6 @@ class FluxEditClient: result_url=result_url, metadata={ "seed": request.seed, - "aspect_ratio": request.aspect_ratio, "safety_tolerance": request.safety_tolerance, "prompt_upsampling": request.prompt_upsampling } diff --git a/src/server/handlers.py b/src/server/handlers.py index 445560e..71515ab 100644 --- a/src/server/handlers.py +++ b/src/server/handlers.py @@ -14,7 +14,6 @@ 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, @@ -155,7 +154,7 @@ class ToolHandlers: 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) + # aspect_ratio is not used in editing (preserve input image aspect ratio) save_to_file = arguments.get('save_to_file', True) logger.info(f"Starting FLUX edit with seed {seed}") @@ -176,7 +175,6 @@ class ToolHandlers: 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 @@ -210,7 +208,6 @@ class ToolHandlers: "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, @@ -236,7 +233,6 @@ class ToolHandlers: 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: @@ -291,7 +287,7 @@ class ToolHandlers: 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) + # aspect_ratio is not used in editing (preserve input image aspect ratio) save_to_file = arguments.get('save_to_file', True) # Check if file exists in input directory @@ -365,7 +361,6 @@ class ToolHandlers: 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 @@ -399,7 +394,6 @@ class ToolHandlers: "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, @@ -428,7 +422,6 @@ class ToolHandlers: 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: @@ -515,91 +508,3 @@ class ToolHandlers: text=f"[ERROR] 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"[ERROR] 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"[ERROR] 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"[SUCCESS] 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"[ERROR] Permission denied: {str(e)}" - )] - except Exception as e: - return [TextContent( - type="text", - text=f"[ERROR] 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"[ERROR] File move error: {str(e)}" - )] diff --git a/src/server/models.py b/src/server/models.py index f298e4e..f554fd2 100644 --- a/src/server/models.py +++ b/src/server/models.py @@ -10,7 +10,6 @@ class ToolName(str, Enum): 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 @@ -56,14 +55,6 @@ TOOL_DEFINITIONS = { 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", @@ -96,14 +87,6 @@ TOOL_DEFINITIONS = { 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", @@ -125,32 +108,6 @@ TOOL_DEFINITIONS = { 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 - ) - ] ) } @@ -180,11 +137,3 @@ class ValidationResult: 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 diff --git a/src/utils/__init__.py b/src/utils/__init__.py index 86fcf36..5def5a0 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -17,7 +17,6 @@ from .image_utils import ( from .validation import ( validate_edit_parameters, validate_file_parameters, - validate_move_file_parameters, validate_image_path_parameter, sanitize_prompt, validate_aspect_ratio_format, @@ -43,7 +42,6 @@ __all__ = [ # Validation utilities 'validate_edit_parameters', 'validate_file_parameters', - 'validate_move_file_parameters', 'validate_image_path_parameter', 'sanitize_prompt', 'validate_aspect_ratio_format', diff --git a/src/utils/validation.py b/src/utils/validation.py index 397a437..63ed364 100644 --- a/src/utils/validation.py +++ b/src/utils/validation.py @@ -69,15 +69,7 @@ def validate_edit_parameters(arguments: Dict[str, Any]) -> Tuple[bool, Optional[ if seed < 0 or seed > 2**32 - 1: return False, "seed must be between 0 and 4294967295" - # Validate optional parameters - if 'aspect_ratio' in arguments: - aspect_ratio = arguments['aspect_ratio'] - if not isinstance(aspect_ratio, str): - return False, "aspect_ratio must be a string" - - valid_ratios = ["1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21"] - if aspect_ratio not in valid_ratios: - return False, f"aspect_ratio must be one of: {', '.join(valid_ratios)}" + # aspect_ratio is not used in editing operations (preserve input image aspect ratio) if 'save_to_file' in arguments: save_to_file = arguments['save_to_file'] @@ -143,15 +135,7 @@ def validate_file_parameters(arguments: Dict[str, Any]) -> Tuple[bool, Optional[ if seed < 0 or seed > 2**32 - 1: return False, "seed must be between 0 and 4294967295" - # Validate optional parameters - if 'aspect_ratio' in arguments: - aspect_ratio = arguments['aspect_ratio'] - if not isinstance(aspect_ratio, str): - return False, "aspect_ratio must be a string" - - valid_ratios = ["1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21"] - if aspect_ratio not in valid_ratios: - return False, f"aspect_ratio must be one of: {', '.join(valid_ratios)}" + # aspect_ratio is not used in editing operations (preserve input image aspect ratio) if 'save_to_file' in arguments: save_to_file = arguments['save_to_file'] @@ -165,52 +149,6 @@ def validate_file_parameters(arguments: Dict[str, Any]) -> Tuple[bool, Optional[ return False, f"Validation error: {str(e)}" -def validate_move_file_parameters(arguments: Dict[str, Any]) -> Tuple[bool, Optional[str]]: - """ - Validate parameters for move_temp_to_output - - Args: - arguments: Tool arguments - - Returns: - tuple: (is_valid, error_message) - """ - try: - # Check required parameters - if 'temp_file_name' not in arguments: - return False, "temp_file_name is required" - - # Validate temp_file_name - temp_file_name = arguments['temp_file_name'] - if not isinstance(temp_file_name, str) or not temp_file_name.strip(): - return False, "temp_file_name must be a non-empty string" - - # Check for path traversal attempts - if '..' in temp_file_name or '/' in temp_file_name or '\\' in temp_file_name: - return False, "temp_file_name cannot contain path separators or '..' for security" - - # Validate optional parameters - if 'output_file_name' in arguments: - output_file_name = arguments['output_file_name'] - if output_file_name is not None: - if not isinstance(output_file_name, str) or not output_file_name.strip(): - return False, "output_file_name must be a non-empty string or None" - - # Check for path traversal attempts - if '..' in output_file_name or '/' in output_file_name or '\\' in output_file_name: - return False, "output_file_name cannot contain path separators or '..' for security" - - if 'copy_only' in arguments: - copy_only = arguments['copy_only'] - if not isinstance(copy_only, bool): - return False, "copy_only must be a boolean" - - return True, None - - except Exception as e: - logger.error(f"Error validating move file parameters: {e}") - return False, f"Validation error: {str(e)}" - def validate_image_path_parameter(arguments: Dict[str, Any]) -> Tuple[bool, Optional[str]]: """