fix bug
This commit is contained in:
@@ -69,15 +69,14 @@ python test_fixes.py # 반드시 실행하여 확인
|
|||||||
"seed": "필수 입력 (일관된 편집 결과를 위해)",
|
"seed": "필수 입력 (일관된 편집 결과를 위해)",
|
||||||
"safety_tolerance": 2,
|
"safety_tolerance": 2,
|
||||||
"output_format": "png",
|
"output_format": "png",
|
||||||
"webhook_url": null,
|
"prompt_upsampling": false
|
||||||
"webhook_secret": null
|
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
#### 2.2 설정 방향
|
#### 2.2 설정 방향
|
||||||
- **prompt**: 토큰 크기 제한이 명시되지 않음 - 제한없이 처리
|
- **prompt**: 토큰 크기 제한이 명시되지 않음 - 제한없이 처리
|
||||||
- **input_image**: 20MB 크기 제한 반영
|
- **input_image**: 20MB 크기 제한 반영
|
||||||
- **aspect_ratio**: 기본값 1:1 또는 16:9 설정 (확인 필요)
|
- **aspect_ratio**: 편집 작업에서는 사용하지 않음 (입력 이미지의 원본 비율 유지)
|
||||||
- **seed**: 반드시 입력받아서 일관된 스타일 유지 및 재현성 보장
|
- **seed**: 반드시 입력받아서 일관된 스타일 유지 및 재현성 보장
|
||||||
- **prompt_upsampling**: false (기본)
|
- **prompt_upsampling**: false (기본)
|
||||||
- **safety_tolerance**: 2 (기본값)
|
- **safety_tolerance**: 2 (기본값)
|
||||||
@@ -135,7 +134,6 @@ D:\Project\little-fairy\flux1-edit\
|
|||||||
"input_image_b64": str, # Base64 인코딩된 입력 이미지
|
"input_image_b64": str, # Base64 인코딩된 입력 이미지
|
||||||
"prompt": str, # 편집 설명
|
"prompt": str, # 편집 설명
|
||||||
"seed": int, # 재현성을 위한 시드값
|
"seed": int, # 재현성을 위한 시드값
|
||||||
"aspect_ratio": str, # "1:1" | "16:9" 등
|
|
||||||
"save_to_file": bool # 파일 저장 여부 (기본: True)
|
"save_to_file": bool # 파일 저장 여부 (기본: True)
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|||||||
26
main.py
26
main.py
@@ -3,7 +3,7 @@
|
|||||||
"""
|
"""
|
||||||
FLUX.1 Edit MCP Server - Fixed Version
|
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
|
- Enhanced error handling and UTF-8 support
|
||||||
- MCP protocol compliance
|
- MCP protocol compliance
|
||||||
- Based on imagen4 server structure
|
- Based on imagen4 server structure
|
||||||
@@ -256,6 +256,28 @@ class FluxEditMCPServer:
|
|||||||
logger.error(f"Error listing tools: {e}", exc_info=True)
|
logger.error(f"Error listing tools: {e}", exc_info=True)
|
||||||
raise
|
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()
|
@self.server.call_tool()
|
||||||
async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[types.TextContent | types.ImageContent]:
|
async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[types.TextContent | types.ImageContent]:
|
||||||
"""Handle tool calls with comprehensive error handling"""
|
"""Handle tool calls with comprehensive error handling"""
|
||||||
@@ -271,8 +293,6 @@ class FluxEditMCPServer:
|
|||||||
return await self.handlers.handle_flux_edit_image_from_file(arguments)
|
return await self.handlers.handle_flux_edit_image_from_file(arguments)
|
||||||
elif name == ToolName.VALIDATE_IMAGE:
|
elif name == ToolName.VALIDATE_IMAGE:
|
||||||
return await self.handlers.handle_validate_image(arguments)
|
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:
|
else:
|
||||||
error_msg = f"Unknown tool: {name}"
|
error_msg = f"Unknown tool: {name}"
|
||||||
logger.error(error_msg)
|
logger.error(error_msg)
|
||||||
|
|||||||
@@ -16,21 +16,22 @@ class Config:
|
|||||||
|
|
||||||
# FLUX.1 Kontext API Configuration
|
# FLUX.1 Kontext API Configuration
|
||||||
API_BASE_URL = "https://api.bfl.ai"
|
API_BASE_URL = "https://api.bfl.ai"
|
||||||
EDIT_ENDPOINT = "/flux-kontext-pro"
|
EDIT_ENDPOINT = "/v1/flux-kontext-pro"
|
||||||
RESULT_ENDPOINT = "/v1/get_result"
|
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"
|
MODEL_NAME = "flux-kontext-pro"
|
||||||
OUTPUT_FORMAT = "png"
|
OUTPUT_FORMAT = "png" # Fixed to PNG format
|
||||||
PROMPT_UPSAMPLING = False
|
PROMPT_UPSAMPLING = False
|
||||||
DEFAULT_SAFETY_TOLERANCE = 2
|
DEFAULT_SAFETY_TOLERANCE = 2
|
||||||
|
|
||||||
# Image size limits
|
# Image size limits
|
||||||
MAX_IMAGE_SIZE_MB = 20 # FLUX.1 Kontext limit
|
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 = [
|
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"
|
DEFAULT_ASPECT_RATIO = "1:1"
|
||||||
|
|
||||||
|
|||||||
@@ -20,7 +20,6 @@ class FluxEditRequest:
|
|||||||
input_image_b64: str
|
input_image_b64: str
|
||||||
prompt: str
|
prompt: str
|
||||||
seed: int
|
seed: int
|
||||||
aspect_ratio: str = "1:1"
|
|
||||||
safety_tolerance: int = 2
|
safety_tolerance: int = 2
|
||||||
output_format: str = "png"
|
output_format: str = "png"
|
||||||
prompt_upsampling: bool = False
|
prompt_upsampling: bool = False
|
||||||
@@ -59,7 +58,12 @@ class FluxEditClient:
|
|||||||
async def _ensure_session(self):
|
async def _ensure_session(self):
|
||||||
"""Ensure aiohttp session is created"""
|
"""Ensure aiohttp session is created"""
|
||||||
if self.session is None or self.session.closed:
|
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)
|
self.session = aiohttp.ClientSession(timeout=timeout)
|
||||||
|
|
||||||
async def close(self):
|
async def close(self):
|
||||||
@@ -71,8 +75,9 @@ class FluxEditClient:
|
|||||||
def _get_headers(self) -> Dict[str, str]:
|
def _get_headers(self) -> Dict[str, str]:
|
||||||
"""Get request headers with API key"""
|
"""Get request headers with API key"""
|
||||||
return {
|
return {
|
||||||
'Content-Type': 'application/json',
|
'accept': 'application/json',
|
||||||
'X-Key': self.config.api_key
|
'x-key': self.config.api_key,
|
||||||
|
'Content-Type': 'application/json'
|
||||||
}
|
}
|
||||||
|
|
||||||
async def _create_edit_request(self, request: FluxEditRequest) -> Optional[str]:
|
async def _create_edit_request(self, request: FluxEditRequest) -> Optional[str]:
|
||||||
@@ -88,54 +93,70 @@ class FluxEditClient:
|
|||||||
try:
|
try:
|
||||||
await self._ensure_session()
|
await self._ensure_session()
|
||||||
|
|
||||||
# Prepare request payload
|
# Prepare request payload for FLUX.1 Kontext API (complete parameters)
|
||||||
payload = {
|
payload = {
|
||||||
"prompt": request.prompt,
|
"prompt": request.prompt,
|
||||||
"input_image": request.input_image_b64,
|
"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,
|
"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
|
# aspect_ratio is not used in editing operations (preserve input image aspect ratio)
|
||||||
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
|
|
||||||
|
|
||||||
url = self.config.get_api_url(self.config.EDIT_ENDPOINT)
|
url = self.config.get_api_url(self.config.EDIT_ENDPOINT)
|
||||||
|
|
||||||
logger.info(f"Creating FLUX edit request to {url}")
|
logger.info(f"Creating FLUX edit request to {url}")
|
||||||
logger.debug(f"Request payload keys: {list(payload.keys())}")
|
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:
|
async with self.session.post(url, json=payload, headers=self._get_headers()) as response:
|
||||||
if response.status == 200:
|
if response.status == 200:
|
||||||
result = await response.json()
|
result = await response.json()
|
||||||
request_id = result.get('id')
|
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}")
|
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:
|
else:
|
||||||
logger.error(f"No request_id in response: {result}")
|
logger.error(f"Missing id or polling_url in response: {result}")
|
||||||
return None
|
return None
|
||||||
else:
|
else:
|
||||||
error_text = await response.text()
|
try:
|
||||||
logger.error(f"Failed to create edit request: {response.status} - {error_text}")
|
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
|
return None
|
||||||
|
|
||||||
except asyncio.TimeoutError:
|
except asyncio.TimeoutError:
|
||||||
logger.error("Timeout creating edit request")
|
logger.error("Timeout creating edit request")
|
||||||
return None
|
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:
|
except Exception as e:
|
||||||
logger.error(f"Error creating edit request: {e}", exc_info=True)
|
logger.error(f"Error creating edit request: {e}", exc_info=True)
|
||||||
return None
|
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:
|
Args:
|
||||||
request_id: Request ID from create_edit_request
|
polling_url: Polling URL from create_edit_request
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
dict: Result data if successful, None otherwise
|
dict: Result data if successful, None otherwise
|
||||||
@@ -143,18 +164,16 @@ class FluxEditClient:
|
|||||||
try:
|
try:
|
||||||
await self._ensure_session()
|
await self._ensure_session()
|
||||||
|
|
||||||
url = self.config.get_api_url(self.config.RESULT_ENDPOINT)
|
# Use the provided polling URL directly
|
||||||
params = {"id": request_id}
|
|
||||||
|
|
||||||
attempts = 0
|
attempts = 0
|
||||||
max_attempts = self.config.max_polling_attempts
|
max_attempts = self.config.max_polling_attempts
|
||||||
interval = self.config.polling_interval
|
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:
|
while attempts < max_attempts:
|
||||||
try:
|
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:
|
if response.status == 200:
|
||||||
result = await response.json()
|
result = await response.json()
|
||||||
|
|
||||||
@@ -260,16 +279,18 @@ class FluxEditClient:
|
|||||||
logger.info(f"Starting FLUX image edit with seed {request.seed}")
|
logger.info(f"Starting FLUX image edit with seed {request.seed}")
|
||||||
|
|
||||||
# Step 1: Create edit request
|
# Step 1: Create edit request
|
||||||
request_id = await self._create_edit_request(request)
|
create_result = await self._create_edit_request(request)
|
||||||
if not request_id:
|
if not create_result:
|
||||||
return FluxEditResponse(
|
return FluxEditResponse(
|
||||||
success=False,
|
success=False,
|
||||||
error_message="Failed to create edit request",
|
error_message="Failed to create edit request",
|
||||||
execution_time=(datetime.now() - start_time).total_seconds()
|
execution_time=(datetime.now() - start_time).total_seconds()
|
||||||
)
|
)
|
||||||
|
|
||||||
|
request_id, polling_url = create_result
|
||||||
|
|
||||||
# Step 2: Poll for result
|
# Step 2: Poll for result
|
||||||
result = await self._poll_result(request_id)
|
result = await self._poll_result(polling_url)
|
||||||
if not result:
|
if not result:
|
||||||
return FluxEditResponse(
|
return FluxEditResponse(
|
||||||
success=False,
|
success=False,
|
||||||
@@ -314,7 +335,6 @@ class FluxEditClient:
|
|||||||
result_url=result_url,
|
result_url=result_url,
|
||||||
metadata={
|
metadata={
|
||||||
"seed": request.seed,
|
"seed": request.seed,
|
||||||
"aspect_ratio": request.aspect_ratio,
|
|
||||||
"safety_tolerance": request.safety_tolerance,
|
"safety_tolerance": request.safety_tolerance,
|
||||||
"prompt_upsampling": request.prompt_upsampling
|
"prompt_upsampling": request.prompt_upsampling
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -14,7 +14,6 @@ from ..utils import (
|
|||||||
validate_edit_parameters,
|
validate_edit_parameters,
|
||||||
validate_file_parameters,
|
validate_file_parameters,
|
||||||
validate_image_path_parameter,
|
validate_image_path_parameter,
|
||||||
validate_move_file_parameters,
|
|
||||||
validate_image_file,
|
validate_image_file,
|
||||||
save_image,
|
save_image,
|
||||||
encode_image_base64,
|
encode_image_base64,
|
||||||
@@ -155,7 +154,7 @@ class ToolHandlers:
|
|||||||
input_image_b64 = arguments['input_image_b64']
|
input_image_b64 = arguments['input_image_b64']
|
||||||
prompt = sanitize_prompt(arguments['prompt'])
|
prompt = sanitize_prompt(arguments['prompt'])
|
||||||
seed = arguments['seed']
|
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)
|
save_to_file = arguments.get('save_to_file', True)
|
||||||
|
|
||||||
logger.info(f"Starting FLUX edit with seed {seed}")
|
logger.info(f"Starting FLUX edit with seed {seed}")
|
||||||
@@ -176,7 +175,6 @@ class ToolHandlers:
|
|||||||
input_image_b64=input_image_b64,
|
input_image_b64=input_image_b64,
|
||||||
prompt=prompt,
|
prompt=prompt,
|
||||||
seed=seed,
|
seed=seed,
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
safety_tolerance=self.config.safety_tolerance,
|
safety_tolerance=self.config.safety_tolerance,
|
||||||
output_format=self.config.OUTPUT_FORMAT,
|
output_format=self.config.OUTPUT_FORMAT,
|
||||||
prompt_upsampling=self.config.prompt_upsampling
|
prompt_upsampling=self.config.prompt_upsampling
|
||||||
@@ -210,7 +208,6 @@ class ToolHandlers:
|
|||||||
"model": self.config.MODEL_NAME,
|
"model": self.config.MODEL_NAME,
|
||||||
"prompt": prompt,
|
"prompt": prompt,
|
||||||
"seed": seed,
|
"seed": seed,
|
||||||
"aspect_ratio": aspect_ratio,
|
|
||||||
"safety_tolerance": self.config.safety_tolerance,
|
"safety_tolerance": self.config.safety_tolerance,
|
||||||
"output_format": self.config.OUTPUT_FORMAT,
|
"output_format": self.config.OUTPUT_FORMAT,
|
||||||
"prompt_upsampling": self.config.prompt_upsampling,
|
"prompt_upsampling": self.config.prompt_upsampling,
|
||||||
@@ -236,7 +233,6 @@ class ToolHandlers:
|
|||||||
text += f"Base name: {base_name}\n"
|
text += f"Base name: {base_name}\n"
|
||||||
if response.image_size:
|
if response.image_size:
|
||||||
text += f"Size: {response.image_size[0]}x{response.image_size[1]}\n"
|
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"
|
text += f"Processing time: {response.execution_time:.1f}s\n"
|
||||||
|
|
||||||
if saved_path:
|
if saved_path:
|
||||||
@@ -291,7 +287,7 @@ class ToolHandlers:
|
|||||||
input_image_name = arguments['input_image_name']
|
input_image_name = arguments['input_image_name']
|
||||||
prompt = sanitize_prompt(arguments['prompt'])
|
prompt = sanitize_prompt(arguments['prompt'])
|
||||||
seed = arguments['seed']
|
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)
|
save_to_file = arguments.get('save_to_file', True)
|
||||||
|
|
||||||
# Check if file exists in input directory
|
# Check if file exists in input directory
|
||||||
@@ -365,7 +361,6 @@ class ToolHandlers:
|
|||||||
input_image_b64=input_image_b64,
|
input_image_b64=input_image_b64,
|
||||||
prompt=prompt,
|
prompt=prompt,
|
||||||
seed=seed,
|
seed=seed,
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
safety_tolerance=self.config.safety_tolerance,
|
safety_tolerance=self.config.safety_tolerance,
|
||||||
output_format=self.config.OUTPUT_FORMAT,
|
output_format=self.config.OUTPUT_FORMAT,
|
||||||
prompt_upsampling=self.config.prompt_upsampling
|
prompt_upsampling=self.config.prompt_upsampling
|
||||||
@@ -399,7 +394,6 @@ class ToolHandlers:
|
|||||||
"model": self.config.MODEL_NAME,
|
"model": self.config.MODEL_NAME,
|
||||||
"prompt": prompt,
|
"prompt": prompt,
|
||||||
"seed": seed,
|
"seed": seed,
|
||||||
"aspect_ratio": aspect_ratio,
|
|
||||||
"safety_tolerance": self.config.safety_tolerance,
|
"safety_tolerance": self.config.safety_tolerance,
|
||||||
"output_format": self.config.OUTPUT_FORMAT,
|
"output_format": self.config.OUTPUT_FORMAT,
|
||||||
"prompt_upsampling": self.config.prompt_upsampling,
|
"prompt_upsampling": self.config.prompt_upsampling,
|
||||||
@@ -428,7 +422,6 @@ class ToolHandlers:
|
|||||||
text += f"Base name: {base_name}\n"
|
text += f"Base name: {base_name}\n"
|
||||||
if response.image_size:
|
if response.image_size:
|
||||||
text += f"Size: {response.image_size[0]}x{response.image_size[1]}\n"
|
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"
|
text += f"Processing time: {response.execution_time:.1f}s\n"
|
||||||
|
|
||||||
if saved_path:
|
if saved_path:
|
||||||
@@ -515,91 +508,3 @@ class ToolHandlers:
|
|||||||
text=f"[ERROR] Validation error: {str(e)}"
|
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)}"
|
|
||||||
)]
|
|
||||||
|
|||||||
@@ -10,7 +10,6 @@ class ToolName(str, Enum):
|
|||||||
FLUX_EDIT_IMAGE = "flux_edit_image"
|
FLUX_EDIT_IMAGE = "flux_edit_image"
|
||||||
FLUX_EDIT_IMAGE_FROM_FILE = "flux_edit_image_from_file"
|
FLUX_EDIT_IMAGE_FROM_FILE = "flux_edit_image_from_file"
|
||||||
VALIDATE_IMAGE = "validate_image"
|
VALIDATE_IMAGE = "validate_image"
|
||||||
MOVE_TEMP_TO_OUTPUT = "move_temp_to_output"
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@@ -56,14 +55,6 @@ TOOL_DEFINITIONS = {
|
|||||||
description="Seed for reproducible results (0 to 4294967295)",
|
description="Seed for reproducible results (0 to 4294967295)",
|
||||||
required=True
|
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(
|
ToolParameter(
|
||||||
name="save_to_file",
|
name="save_to_file",
|
||||||
type="boolean",
|
type="boolean",
|
||||||
@@ -96,14 +87,6 @@ TOOL_DEFINITIONS = {
|
|||||||
description="Seed for reproducible results (0 to 4294967295)",
|
description="Seed for reproducible results (0 to 4294967295)",
|
||||||
required=True
|
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(
|
ToolParameter(
|
||||||
name="save_to_file",
|
name="save_to_file",
|
||||||
type="boolean",
|
type="boolean",
|
||||||
@@ -125,32 +108,6 @@ TOOL_DEFINITIONS = {
|
|||||||
required=True
|
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
|
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
|
|
||||||
|
|||||||
@@ -17,7 +17,6 @@ from .image_utils import (
|
|||||||
from .validation import (
|
from .validation import (
|
||||||
validate_edit_parameters,
|
validate_edit_parameters,
|
||||||
validate_file_parameters,
|
validate_file_parameters,
|
||||||
validate_move_file_parameters,
|
|
||||||
validate_image_path_parameter,
|
validate_image_path_parameter,
|
||||||
sanitize_prompt,
|
sanitize_prompt,
|
||||||
validate_aspect_ratio_format,
|
validate_aspect_ratio_format,
|
||||||
@@ -43,7 +42,6 @@ __all__ = [
|
|||||||
# Validation utilities
|
# Validation utilities
|
||||||
'validate_edit_parameters',
|
'validate_edit_parameters',
|
||||||
'validate_file_parameters',
|
'validate_file_parameters',
|
||||||
'validate_move_file_parameters',
|
|
||||||
'validate_image_path_parameter',
|
'validate_image_path_parameter',
|
||||||
'sanitize_prompt',
|
'sanitize_prompt',
|
||||||
'validate_aspect_ratio_format',
|
'validate_aspect_ratio_format',
|
||||||
|
|||||||
@@ -69,15 +69,7 @@ def validate_edit_parameters(arguments: Dict[str, Any]) -> Tuple[bool, Optional[
|
|||||||
if seed < 0 or seed > 2**32 - 1:
|
if seed < 0 or seed > 2**32 - 1:
|
||||||
return False, "seed must be between 0 and 4294967295"
|
return False, "seed must be between 0 and 4294967295"
|
||||||
|
|
||||||
# Validate optional parameters
|
# aspect_ratio is not used in editing operations (preserve input image aspect ratio)
|
||||||
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)}"
|
|
||||||
|
|
||||||
if 'save_to_file' in arguments:
|
if 'save_to_file' in arguments:
|
||||||
save_to_file = arguments['save_to_file']
|
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:
|
if seed < 0 or seed > 2**32 - 1:
|
||||||
return False, "seed must be between 0 and 4294967295"
|
return False, "seed must be between 0 and 4294967295"
|
||||||
|
|
||||||
# Validate optional parameters
|
# aspect_ratio is not used in editing operations (preserve input image aspect ratio)
|
||||||
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)}"
|
|
||||||
|
|
||||||
if 'save_to_file' in arguments:
|
if 'save_to_file' in arguments:
|
||||||
save_to_file = arguments['save_to_file']
|
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)}"
|
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]]:
|
def validate_image_path_parameter(arguments: Dict[str, Any]) -> Tuple[bool, Optional[str]]:
|
||||||
"""
|
"""
|
||||||
|
|||||||
Reference in New Issue
Block a user