From f1f5ad04f1b4011cc37ffa9343ff49a23e6801f0 Mon Sep 17 00:00:00 2001 From: ened Date: Tue, 26 Aug 2025 02:35:44 +0900 Subject: [PATCH] initial implementation --- .env.example | 26 ++ .gitignore | 100 ++++++ CLAUDE.md | 261 +++++++++++++++ README.md | 198 +++++++++++ check_dependencies.py | 150 +++++++++ install_dependencies.bat | 122 +++++++ install_dependencies.ps1 | 138 ++++++++ main.py | 123 +++++++ requirements.txt | 29 ++ run.bat | 74 +++++ run.ps1 | 83 +++++ src/__init__.py | 5 + src/connector/__init__.py | 33 ++ src/connector/config.py | 272 +++++++++++++++ src/connector/flux_client.py | 334 +++++++++++++++++++ src/server/__init__.py | 7 + src/server/handlers.py | 605 ++++++++++++++++++++++++++++++++++ src/server/handlers_backup.py | 605 ++++++++++++++++++++++++++++++++++ src/server/mcp_server.py | 178 ++++++++++ src/server/models.py | 190 +++++++++++ src/utils/__init__.py | 55 ++++ src/utils/image_utils.py | 385 ++++++++++++++++++++++ src/utils/validation.py | 387 ++++++++++++++++++++++ start.bat | 25 ++ test_fixes.py | 119 +++++++ tests/__init__.py | 0 tests/run_tests.py | 106 ++++++ tests/test_config.py | 159 +++++++++ tests/test_flux_client.py | 308 +++++++++++++++++ tests/test_handlers.py | 360 ++++++++++++++++++++ tests/test_image_utils.py | 197 +++++++++++ tests/test_validation.py | 218 ++++++++++++ troubleshoot.bat | 150 +++++++++ 33 files changed, 6002 insertions(+) create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 CLAUDE.md create mode 100644 README.md create mode 100644 check_dependencies.py create mode 100644 install_dependencies.bat create mode 100644 install_dependencies.ps1 create mode 100644 main.py create mode 100644 requirements.txt create mode 100644 run.bat create mode 100644 run.ps1 create mode 100644 src/__init__.py create mode 100644 src/connector/__init__.py create mode 100644 src/connector/config.py create mode 100644 src/connector/flux_client.py create mode 100644 src/server/__init__.py create mode 100644 src/server/handlers.py create mode 100644 src/server/handlers_backup.py create mode 100644 src/server/mcp_server.py create mode 100644 src/server/models.py create mode 100644 src/utils/__init__.py create mode 100644 src/utils/image_utils.py create mode 100644 src/utils/validation.py create mode 100644 start.bat create mode 100644 test_fixes.py create mode 100644 tests/__init__.py create mode 100644 tests/run_tests.py create mode 100644 tests/test_config.py create mode 100644 tests/test_flux_client.py create mode 100644 tests/test_handlers.py create mode 100644 tests/test_image_utils.py create mode 100644 tests/test_validation.py create mode 100644 troubleshoot.bat diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..0f5927d --- /dev/null +++ b/.env.example @@ -0,0 +1,26 @@ +# Environment variables for FLUX.1 Edit MCP Server + +# FLUX API Configuration +FLUX_API_KEY=your_flux_api_key_here +FLUX_API_BASE_URL=https://api.bfl.ai + +# Server Configuration +LOG_LEVEL=INFO +MAX_IMAGE_SIZE_MB=20 +DEFAULT_TIMEOUT=300 +POLLING_INTERVAL=2 +MAX_POLLING_ATTEMPTS=150 + +# Path Configuration +INPUT_PATH=./input_images +GENERATED_IMAGES_PATH=./generated_images + +# File Management +OUTPUT_FILENAME_PREFIX=fluxedit +SAVE_PARAMETERS=true + +# Default Settings +DEFAULT_ASPECT_RATIO=1:1 +DEFAULT_SAFETY_TOLERANCE=2 +OUTPUT_FORMAT=png +PROMPT_UPSAMPLING=false diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7c5d94b --- /dev/null +++ b/.gitignore @@ -0,0 +1,100 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Virtual environments +.env +.env.local +.env.development +.env.test +.env.production +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# IDEs +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS generated files +.DS_Store +.DS_Store? +._* +.Spotlight-V100 +.Trashes +ehthumbs.db +Thumbs.db + +# Project specific +flux1-edit.log +*.log +temp/ +.temp/ +generated_images/*.png +generated_images/*.jpg +generated_images/*.jpeg +generated_images/*.webp +generated_images/*.json +!generated_images/.gitkeep + +# Keep input_images directory but ignore contents +input_images/* +!input_images/.gitkeep + +# API keys and secrets +.env.local +.env.production +config/secrets.json + +# Temporary files +*.tmp +*.bak diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..e622efd --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,261 @@ +# FLUX.1 Kontext MCP Server 개발 가이드 + +## 프로젝트 개요 + +**프로젝트명**: flux1-edit +**목적**: FLUX.1 Kontext 모델 API를 이용한 MCP server 및 connector 구현 +**특화 기능**: 이미지 생성이 아닌 편집(editing)에 특화된 MCP 도구 제공 + +## ⚠️ 소스 코드 작성 규칙 (중요!) + +### MCP 서버 호환성을 위한 필수 규칙 + +#### 1. Unicode 문자 사용 금지 +```python +# 금지: ✅❌🎲📁📐💾 등 모든 이모지/특수문자 +# 권장: [SUCCESS] [ERROR] [INFO] [OK] [FAILED] + +print("[SUCCESS] 작업 완료") # ✅ +print("✅ 작업 완료") # ❌ +``` + +#### 2. 콘솔 출력 금지 (stdout 사용 안함) +```python +# 금지: print() - JSON 파싱 오류 유발 +# 권장: logger 사용 + +logger.info("서버 시작") # ✅ +print("서버 시작") # ❌ +``` + +#### 3. 배치 파일에 UTF-8 설정 필수 +```batch +chcp 65001 >nul 2>&1 +set PYTHONIOENCODING=utf-8 +set PYTHONUTF8=1 +python main.py +``` + +#### 4. 로깅 설정 +```python +logging.basicConfig( + handlers=[ + logging.FileHandler('app.log', encoding='utf-8') # 파일만 + # logging.StreamHandler(sys.stdout) # 금지! + ] +) +``` + +### 수정 후 검증 +```bash +python test_fixes.py # 반드시 실행하여 확인 +``` + +--- + +## API 스펙 분석 (FLUX.1 Kontext) + +### 1. 엔드포인트 +- **생성 요청**: `POST /flux-kontext-pro` +- **결과 조회**: `GET /v1/get_result?id={request_id}` + +### 2. 주요 파라미터 설정 + +#### 2.1 입력 파라미터 +```json +{ + "prompt": "편집 설명 (텍스트)", + "input_image": "base64 인코딩된 이미지 (20MB 제한)", + "seed": "필수 입력 (일관된 편집 결과를 위해)", + "safety_tolerance": 2, + "output_format": "png", + "webhook_url": null, + "webhook_secret": null +} +``` + +#### 2.2 설정 방향 +- **prompt**: 토큰 크기 제한이 명시되지 않음 - 제한없이 처리 +- **input_image**: 20MB 크기 제한 반영 +- **aspect_ratio**: 기본값 1:1 또는 16:9 설정 (확인 필요) +- **seed**: 반드시 입력받아서 일관된 스타일 유지 및 재현성 보장 +- **prompt_upsampling**: false (기본) +- **safety_tolerance**: 2 (기본값) +- **output_format**: png 고정 +- **webhook_url, webhook_secret**: 사용 안 함 + +### 3. 워크플로우 +1. 요청 생성 → `request_id` 반환 +2. `request_id`로 결과 폴링 +3. `result['sample']`에서 서명된 URL 획득 +4. URL은 10분간 유효 (즉시 다운로드 필요) + +## 개발 구조 + +### 1. 디렉터리 구조 +``` +D:\Project\little-fairy\flux1-edit\ +├── CLAUDE.md # 이 가이드 문서 +├── README.md # 프로젝트 설명서 +├── requirements.txt # Python 의존성 +├── main.py # 메인 서버 실행 파일 +├── .env # 환경변수 (FLUX API 키) +├── .env.example # 환경변수 예시 +├── input_images/ # 입력 이미지 디렉터리 +├── generated_images/ # 출력 이미지 디렉터리 +├── src/ +│ ├── __init__.py +│ ├── server/ +│ │ ├── __init__.py +│ │ ├── mcp_server.py # MCP 서버 메인 +│ │ ├── handlers.py # 도구 핸들러들 +│ │ └── models.py # 데이터 모델들 +│ ├── connector/ +│ │ ├── __init__.py +│ │ ├── config.py # 설정 관리 +│ │ └── flux_client.py # FLUX API 클라이언트 +│ └── utils/ +│ ├── __init__.py +│ ├── image_utils.py # 이미지 처리 유틸 +│ └── validation.py # 입력 검증 +└── tests/ # 단위 테스트 + ├── __init__.py + ├── test_config.py + ├── test_flux_client.py + └── test_handlers.py +``` + +### 2. MCP 도구 정의 + +#### 2.1 flux_edit_image (메인 도구) +- **설명**: FLUX.1 Kontext로 이미지 편집 +- **입력**: + ```python + { + "input_image_b64": str, # Base64 인코딩된 입력 이미지 + "prompt": str, # 편집 설명 + "seed": int, # 재현성을 위한 시드값 + "aspect_ratio": str, # "1:1" | "16:9" 등 + "save_to_file": bool # 파일 저장 여부 (기본: True) + } + ``` +- **출력**: 편집된 이미지 + 메타데이터 + +#### 2.2 지원 도구들 (gpt-edit 참고) +- `validate_image`: 이미지 검증 +- `flux_edit_image_from_file`: 파일에서 읽어서 편집 +- `move_temp_to_output`: 임시 파일을 출력 디렉터리로 이동 + +### 3. 주요 구현 포인트 + +#### 3.1 Config 클래스 (src/connector/config.py) +```python +class Config: + # FLUX.1 Kontext 전용 설정 + API_BASE_URL = "https://api.bfl.ai" + MODEL_NAME = "flux-kontext-pro" + MAX_IMAGE_SIZE_MB = 20 # FLUX 제한에 맞춤 + DEFAULT_ASPECT_RATIO = "1:1" + DEFAULT_SAFETY_TOLERANCE = 2 + OUTPUT_FORMAT = "png" + PROMPT_UPSAMPLING = False + + # 환경변수 + FLUX_API_KEY = os.getenv('FLUX_API_KEY') + + # 파일 경로 + input_path: Path + generated_images_path: Path +``` + +#### 3.2 FLUX Client (src/connector/flux_client.py) +```python +class FluxEditClient: + async def edit_image(self, request: FluxEditRequest) -> FluxEditResponse: + # 1. 편집 요청 생성 + # 2. 결과 폴링 + # 3. 이미지 다운로드 + # 4. 응답 생성 + pass + +class FluxEditRequest: + input_image_b64: str + prompt: str + seed: int + aspect_ratio: str = "1:1" + +class FluxEditResponse: + success: bool + edited_image_data: bytes + image_size: Tuple[int, int] + execution_time: float + error_message: Optional[str] +``` + +#### 3.3 핸들러 구현 (src/server/handlers.py) +- gpt-edit의 `handle_edit_image_from_file` 구조 참고 +- FLUX API 특성에 맞게 수정: + - 마스크 기능 없음 (FLUX.1 Kontext는 프롬프트만으로 편집) + - 폴링 방식의 비동기 처리 + - 20MB 이미지 크기 제한 + +## 개발 우선순위 + +### Phase 1: 핵심 기능 구현 +1. ✅ 프로젝트 구조 생성 +2. ⏳ Config 및 환경설정 +3. ⏳ FLUX API 클라이언트 구현 +4. ⏳ 기본 MCP 서버 구현 +5. ⏳ `flux_edit_image` 도구 구현 + +### Phase 2: 편의 기능 추가 +1. 파일 기반 편집 (`flux_edit_image_from_file`) +2. 이미지 검증 도구 +3. 결과 파일 관리 도구 + +### Phase 3: 테스트 및 최적화 +1. 단위 테스트 작성 +2. 에러 핸들링 강화 +3. 성능 최적화 + +## 참고사항 + +### 1. gpt-edit와의 차이점 +- **API 구조**: OpenAI는 동기식, FLUX는 비동기식 (폴링) +- **마스크 기능**: OpenAI는 마스크 지원, FLUX는 프롬프트만 +- **이미지 크기**: OpenAI 4MB → FLUX 20MB +- **출력 형식**: 다양한 형식 → PNG 고정 + +### 2. 보안 고려사항 +- FLUX API 키 환경변수 관리 +- 임시 파일 자동 정리 +- 이미지 데이터 메모리 관리 + +### 3. 에러 핸들링 +- 네트워크 타임아웃 +- API 할당량 초과 +- 이미지 형식/크기 오류 +- 폴링 중 서버 오류 + +## 환경변수 설정 + +```bash +# .env 파일 예시 +FLUX_API_KEY=your_flux_api_key_here +LOG_LEVEL=INFO +MAX_IMAGE_SIZE_MB=20 +DEFAULT_TIMEOUT=300 +INPUT_PATH=./input_images +GENERATED_IMAGES_PATH=./generated_images +SAVE_PARAMETERS=true +``` + +## 다음 단계 + +1. **환경 설정 파일 생성** (.env, requirements.txt) +2. **기본 구조 코드 작성** (config.py, flux_client.py) +3. **MCP 서버 메인 로직** (mcp_server.py, handlers.py) +4. **단위 테스트 작성 및 실행** +5. **Claude에서 실제 테스트** + +이 가이드를 바탕으로 step-by-step으로 구현을 진행하면 됩니다. diff --git a/README.md b/README.md new file mode 100644 index 0000000..9beb782 --- /dev/null +++ b/README.md @@ -0,0 +1,198 @@ +# FLUX.1 Edit MCP Server + +FLUX.1 Kontext 모델을 이용한 이미지 편집 전용 MCP (Model Context Protocol) 서버입니다. + +## 주요 기능 + +- **이미지 편집**: FLUX.1 Kontext API를 이용한 고품질 이미지 편집 +- **시드 기반 재현성**: 동일한 시드로 일관된 편집 결과 보장 +- **대용량 이미지 지원**: 최대 20MB 이미지 처리 +- **MCP 통합**: Claude와 완전 통합된 이미지 편집 도구 + +## 설치 및 설정 + +### 빠른 설치 (권장) + +**Windows 배치 파일 사용:** +```bash +# 의존성 자동 설치 +install_dependencies.bat + +# 서버 실행 +run.bat +``` + +**PowerShell 스크립트 사용:** +```powershell +# 의존성 자동 설치 +.\install_dependencies.ps1 + +# 서버 실행 +.\run.ps1 +``` + +### 수동 설치 + +1. **Python 가상 환경 생성** +```bash +python -m venv venv +venv\Scripts\activate # Windows +source venv/bin/activate # Linux/Mac +``` + +2. **의존성 설치** +```bash +pip install -r requirements.txt +``` + +3. **환경변수 설정** +```bash +copy .env.example .env # Windows +cp .env.example .env # Linux/Mac +# .env 파일에서 FLUX_API_KEY 설정 +``` + +4. **필요한 디렉터리 생성** +```bash +mkdir input_images generated_images temp +``` + +5. **서버 실행** +```bash +python main.py +``` + +## 문제 해결 + +### 일반적인 오류들 + +#### "ModuleNotFoundError: No module named 'aiohttp'" +```bash +# 문제 진단 +troubleshoot.bat + +# 또는 직접 설치 +pip install aiohttp==3.11.7 +``` + +#### 가상 환경 활성화 문제 +```bash +# 가상 환경 재생성 +rmdir /s venv +python -m venv venv +venv\Scripts\activate +pip install -r requirements.txt +``` + +#### PowerShell 실행 정책 오류 +```powershell +Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser +``` + +### 문제 진단 도구 + +**자동 진단 실행:** +```bash +troubleshoot.bat +``` + +이 스크립트는 다음을 확인합니다: +- Python 설치 +- pip 가용성 +- 가상 환경 상태 +- 필수 패키지 설치 상태 +- 설정 파일 존재 +- 로컬 모듈 import 가능성 + +### 수동 의존성 확인 + +```bash +python -c " +import aiohttp +import httpx +import mcp +from PIL import Image +import dotenv +import pydantic +print('All dependencies are working!') +" +``` + +## 환경 요구사항 + +- Python 3.8 이상 +- 안정적인 인터넷 연결 (FLUX API 통신용) +- FLUX API 키 + +## 사용 가능한 도구 + +### flux_edit_image +FLUX.1 Kontext로 이미지를 편집합니다. + +**파라미터:** +- `input_image_b64`: Base64 인코딩된 입력 이미지 +- `prompt`: 편집 설명 +- `seed`: 재현성을 위한 시드값 +- `aspect_ratio`: "1:1" | "16:9" 등 (기본: "1:1") +- `save_to_file`: 파일 저장 여부 (기본: True) + +### validate_image +이미지 파일의 유효성을 검증합니다. + +### move_temp_to_output +임시 파일을 출력 디렉터리로 이동합니다. + +## 디렉터리 구조 + +``` +flux1-edit/ +├── main.py # 메인 실행 파일 +├── requirements.txt # Python 의존성 +├── .env # 환경 설정 (FLUX_API_KEY) +├── .env.example # 환경 설정 예제 +├── install_dependencies.bat # Windows 자동 설치 +├── install_dependencies.ps1 # PowerShell 자동 설치 +├── run.bat # Windows 실행 스크립트 +├── run.ps1 # PowerShell 실행 스크립트 +├── troubleshoot.bat # 문제 진단 도구 +├── src/ # 소스 코드 +│ ├── connector/ # FLUX API 연결 +│ ├── server/ # MCP 서버 +│ └── utils/ # 유틸리티 +├── input_images/ # 입력 이미지 +├── generated_images/ # 출력 이미지 +├── temp/ # 임시 파일 +└── tests/ # 테스트 파일 +``` + +## Claude 설정 + +MCP 서버를 Claude에 연결하려면: + +1. Claude 설정에서 MCP 서버 추가 +2. 서버 명령: `python D:\Project\flux1-edit\main.py` +3. 작업 디렉터리: `D:\Project\flux1-edit` + +## API 제한사항 + +- 최대 이미지 크기: 20MB +- 출력 형식: PNG 고정 +- 결과 URL 유효시간: 10분 +- 편집 타임아웃: 5분 + +## 로그 및 디버깅 + +- 로그 파일: `flux1-edit.log` +- 로그 레벨: INFO (콘솔), DEBUG (파일) +- 오류 추적: 전체 스택 트레이스 포함 + +## 지원 및 문의 + +문제가 지속되는 경우: +1. `troubleshoot.bat` 실행하여 진단 정보 확인 +2. `flux1-edit.log` 파일에서 자세한 오류 로그 확인 +3. 가상 환경 재생성 후 재시도 + +## 라이선스 + +MIT License diff --git a/check_dependencies.py b/check_dependencies.py new file mode 100644 index 0000000..1e3548e --- /dev/null +++ b/check_dependencies.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +""" +FLUX.1 Edit MCP Server - Dependency Check Script +This script checks if all required dependencies are properly installed +""" + +import sys +import importlib.util +from typing import List, Tuple + +def check_dependency(module_name: str, package_name: str = None) -> Tuple[bool, str]: + """ + Check if a dependency is installed and importable + + Args: + module_name: Name of the module to import + package_name: Name of the package to install (if different from module) + + Returns: + Tuple of (success, message) + """ + if package_name is None: + package_name = module_name + + try: + spec = importlib.util.find_spec(module_name) + if spec is None: + return False, f"[MISSING] {module_name} not found - install with: pip install {package_name}" + + # Try to actually import it + module = importlib.import_module(module_name) + + # Get version if available + version = getattr(module, '__version__', 'unknown') + return True, f"[OK] {module_name} {version}" + + except ImportError as e: + return False, f"[ERROR] {module_name} import failed: {e}" + except Exception as e: + return False, f"[ERROR] {module_name} error: {e}" + +def check_local_modules() -> List[Tuple[bool, str]]: + """Check local project modules""" + results = [] + + try: + # Add src to path temporarily + import os + from pathlib import Path + + src_path = Path(__file__).parent / 'src' + if src_path.exists(): + sys.path.insert(0, str(src_path)) + + # Test local imports + try: + from connector.config import Config + results.append((True, "[OK] Local Config module")) + except Exception as e: + results.append((False, f"[ERROR] Local Config module: {e}")) + + try: + from connector.flux_client import FluxEditClient + results.append((True, "[OK] Local FluxEditClient module")) + except Exception as e: + results.append((False, f"[ERROR] Local FluxEditClient module: {e}")) + + try: + from server.mcp_server import FluxEditMCPServer + results.append((True, "[OK] Local MCP Server module")) + except Exception as e: + results.append((False, f"[ERROR] Local MCP Server module: {e}")) + + except Exception as e: + results.append((False, f"[ERROR] Local module check failed: {e}")) + + return results + +def main(): + """Main dependency check function""" + print("FLUX.1 Edit MCP Server - Dependency Check") + print("=========================================") + print(f"Python version: {sys.version}") + print(f"Python executable: {sys.executable}") + print() + + # Required dependencies with their install names + dependencies = [ + ("aiohttp", "aiohttp==3.11.7"), + ("httpx", "httpx==0.28.1"), + ("mcp", "mcp==1.1.0"), + ("PIL", "Pillow==11.0.0"), + ("dotenv", "python-dotenv==1.0.1"), + ("pydantic", "pydantic==2.10.3"), + ("structlog", "structlog==24.4.0"), + ] + + # Optional dependencies for development + optional_dependencies = [ + ("pytest", "pytest==8.3.4"), + ("black", "black==24.10.0"), + ] + + all_good = True + + print("Checking required dependencies...") + print("-" * 50) + + for module_name, package_name in dependencies: + success, message = check_dependency(module_name, package_name) + print(message) + if not success: + all_good = False + + print() + print("Checking optional dependencies...") + print("-" * 50) + + for module_name, package_name in optional_dependencies: + success, message = check_dependency(module_name, package_name) + print(message) + + print() + print("Checking local modules...") + print("-" * 50) + + local_results = check_local_modules() + for success, message in local_results: + print(message) + if not success: + all_good = False + + print() + print("=" * 50) + + if all_good: + print("[SUCCESS] All required dependencies are installed and working!") + print("You can now run: python main.py") + return 0 + else: + print("[FAILED] Some dependencies are missing or broken.") + print() + print("To fix this, try:") + print("1. Run: install_dependencies.bat (Windows)") + print("2. Or run: pip install -r requirements.txt") + print("3. Or run individual pip install commands shown above") + return 1 + +if __name__ == "__main__": + sys.exit(main()) diff --git a/install_dependencies.bat b/install_dependencies.bat new file mode 100644 index 0000000..90b0b97 --- /dev/null +++ b/install_dependencies.bat @@ -0,0 +1,122 @@ +@echo off +echo Installing FLUX.1 Edit MCP Server Dependencies... +echo =============================================== + +REM Check if Python is available +python --version >nul 2>&1 +if errorlevel 1 ( + echo Error: Python is not installed or not in PATH + echo Please install Python 3.8 or higher from https://python.org + pause + exit /b 1 +) + +echo Python version: +python --version +echo. + +REM Check if pip is available +pip --version >nul 2>&1 +if errorlevel 1 ( + echo Error: pip is not available + echo Please ensure pip is installed with Python + pause + exit /b 1 +) + +REM Update pip to latest version +echo Updating pip... +python -m pip install --upgrade pip +if errorlevel 1 ( + echo Warning: Failed to upgrade pip, continuing with current version +) + +REM Check if virtual environment exists, create if not +if not exist "venv" ( + echo Creating virtual environment... + python -m venv venv + if errorlevel 1 ( + echo Error: Failed to create virtual environment + echo This might be due to permissions or Python installation issues + pause + exit /b 1 + ) + echo Virtual environment created successfully. +) else ( + echo Virtual environment already exists. +) + +REM Activate virtual environment +echo Activating virtual environment... +call venv\Scripts\activate.bat +if errorlevel 1 ( + echo Error: Failed to activate virtual environment + echo Trying to recreate virtual environment... + rmdir /s /q venv + python -m venv venv + call venv\Scripts\activate.bat + if errorlevel 1 ( + echo Error: Still failed to activate virtual environment + pause + exit /b 1 + ) +) + +echo Virtual environment activated. +echo Current Python location: +where python +echo. + +REM Install dependencies with explicit error handling +echo Installing dependencies from requirements.txt... +pip install -r requirements.txt --no-cache-dir -v +if errorlevel 1 ( + echo Error: Failed to install some dependencies + echo Trying to install critical dependencies individually... + echo. + + echo Installing aiohttp... + pip install aiohttp==3.11.7 + + echo Installing httpx... + pip install httpx==0.28.1 + + echo Installing mcp... + pip install mcp==1.1.0 + + echo Installing Pillow... + pip install Pillow==11.0.0 + + echo Installing python-dotenv... + pip install python-dotenv==1.0.1 + + echo Installing pydantic... + pip install pydantic==2.10.3 + + echo Installing structlog... + pip install structlog==24.4.0 +) + +echo. +echo Verifying installations... +python -c "import aiohttp; print(f'aiohttp version: {aiohttp.__version__}')" 2>nul +if errorlevel 1 ( + echo Error: aiohttp is not properly installed + echo Trying alternative installation method... + pip install --force-reinstall aiohttp +) + +python -c "import mcp; print('mcp module imported successfully')" 2>nul +if errorlevel 1 ( + echo Warning: mcp module check failed, but this might be normal +) + +echo. +echo Installation completed! +echo. +echo Next steps: +echo 1. Copy .env.example to .env: copy .env.example .env +echo 2. Edit .env file and add your FLUX_API_KEY +echo 3. Run the server with: run.bat or python main.py +echo. +pause diff --git a/install_dependencies.ps1 b/install_dependencies.ps1 new file mode 100644 index 0000000..2b59663 --- /dev/null +++ b/install_dependencies.ps1 @@ -0,0 +1,138 @@ +# PowerShell script to install FLUX.1 Edit MCP Server Dependencies + +Write-Host "Installing FLUX.1 Edit MCP Server Dependencies..." -ForegroundColor Green +Write-Host "===============================================" -ForegroundColor Green + +# Check if Python is available +try { + $pythonVersion = python --version + Write-Host "Found Python: $pythonVersion" -ForegroundColor Blue +} catch { + Write-Host "Error: Python is not installed or not in PATH" -ForegroundColor Red + Write-Host "Please install Python 3.8 or higher from https://python.org" -ForegroundColor Yellow + Read-Host "Press Enter to exit" + exit 1 +} + +# Check if pip is available +try { + $pipVersion = pip --version + Write-Host "Found pip: $pipVersion" -ForegroundColor Blue +} catch { + Write-Host "Error: pip is not available" -ForegroundColor Red + Write-Host "Please ensure pip is installed with Python" -ForegroundColor Yellow + Read-Host "Press Enter to exit" + exit 1 +} + +# Update pip to latest version +Write-Host "Updating pip..." -ForegroundColor Blue +try { + python -m pip install --upgrade pip + Write-Host "pip updated successfully" -ForegroundColor Green +} catch { + Write-Host "Warning: Failed to upgrade pip, continuing with current version" -ForegroundColor Yellow +} + +# Check if virtual environment exists, create if not +if (-not (Test-Path "venv")) { + Write-Host "Creating virtual environment..." -ForegroundColor Yellow + python -m venv venv + if ($LASTEXITCODE -ne 0) { + Write-Host "Error: Failed to create virtual environment" -ForegroundColor Red + Write-Host "This might be due to permissions or Python installation issues" -ForegroundColor Yellow + Read-Host "Press Enter to exit" + exit 1 + } + Write-Host "Virtual environment created successfully." -ForegroundColor Green +} else { + Write-Host "Virtual environment already exists." -ForegroundColor Blue +} + +# Activate virtual environment +Write-Host "Activating virtual environment..." -ForegroundColor Blue +try { + & "venv\Scripts\Activate.ps1" + if ($LASTEXITCODE -ne 0) { + throw "Activation failed" + } + Write-Host "Virtual environment activated." -ForegroundColor Green +} catch { + Write-Host "Error: Failed to activate virtual environment" -ForegroundColor Red + Write-Host "Trying to recreate virtual environment..." -ForegroundColor Yellow + Remove-Item -Recurse -Force "venv" -ErrorAction SilentlyContinue + python -m venv venv + & "venv\Scripts\Activate.ps1" + if ($LASTEXITCODE -ne 0) { + Write-Host "Error: Still failed to activate virtual environment" -ForegroundColor Red + Write-Host "You might need to enable PowerShell script execution:" -ForegroundColor Yellow + Write-Host "Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser" -ForegroundColor Yellow + Read-Host "Press Enter to exit" + exit 1 + } +} + +Write-Host "Current Python location:" -ForegroundColor Blue +& where.exe python + +# Install dependencies with explicit error handling +Write-Host "`nInstalling dependencies from requirements.txt..." -ForegroundColor Blue +try { + pip install -r requirements.txt --no-cache-dir -v + if ($LASTEXITCODE -ne 0) { + throw "pip install failed" + } + Write-Host "Dependencies installed successfully!" -ForegroundColor Green +} catch { + Write-Host "Error: Failed to install some dependencies" -ForegroundColor Red + Write-Host "Trying to install critical dependencies individually..." -ForegroundColor Yellow + + $criticalPackages = @( + "aiohttp==3.11.7", + "httpx==0.28.1", + "mcp==1.1.0", + "Pillow==11.0.0", + "python-dotenv==1.0.1", + "pydantic==2.10.3", + "structlog==24.4.0" + ) + + foreach ($package in $criticalPackages) { + Write-Host "Installing $package..." -ForegroundColor Blue + pip install $package + } +} + +# Verify installations +Write-Host "`nVerifying installations..." -ForegroundColor Blue +try { + $aiohttpVersion = python -c "import aiohttp; print(f'aiohttp version: {aiohttp.__version__}')" 2>$null + if ($LASTEXITCODE -eq 0) { + Write-Host $aiohttpVersion -ForegroundColor Green + } else { + throw "aiohttp import failed" + } +} catch { + Write-Host "Error: aiohttp is not properly installed" -ForegroundColor Red + Write-Host "Trying alternative installation method..." -ForegroundColor Yellow + pip install --force-reinstall aiohttp +} + +try { + python -c "import mcp; print('mcp module imported successfully')" 2>$null + if ($LASTEXITCODE -eq 0) { + Write-Host "mcp module imported successfully" -ForegroundColor Green + } else { + Write-Host "Warning: mcp module check failed, but this might be normal" -ForegroundColor Yellow + } +} catch { + Write-Host "Warning: mcp module verification failed" -ForegroundColor Yellow +} + +Write-Host "`nInstallation completed!" -ForegroundColor Green +Write-Host "`nNext steps:" -ForegroundColor Yellow +Write-Host "1. Copy .env.example to .env: Copy-Item .env.example .env" -ForegroundColor White +Write-Host "2. Edit .env file and add your FLUX_API_KEY" -ForegroundColor White +Write-Host "3. Run the server with: .\run.ps1 or python main.py" -ForegroundColor White +Write-Host "" +Read-Host "Press Enter to exit" diff --git a/main.py b/main.py new file mode 100644 index 0000000..38958d5 --- /dev/null +++ b/main.py @@ -0,0 +1,123 @@ +"""Main entry point for FLUX.1 Edit MCP Server""" + +import asyncio +import logging +import sys +import traceback +from pathlib import Path + +# Add src to path for imports +sys.path.insert(0, str(Path(__file__).parent / 'src')) + + +def setup_logging(): + """Setup logging configuration""" + # Only log to file, not stdout to avoid JSON parsing issues + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('flux1-edit.log', mode='a', encoding='utf-8') + ] + ) + + # Set specific loggers + logging.getLogger('aiohttp').setLevel(logging.WARNING) + logging.getLogger('PIL').setLevel(logging.WARNING) + + +def check_dependencies(): + """Check if all required dependencies are available - silently check for MCP compatibility""" + missing_deps = [] + + try: + import aiohttp + # Silent check - no print to avoid JSON parsing issues + except ImportError as e: + missing_deps.append(f"aiohttp: {e}") + + try: + import httpx + # Silent check - no print to avoid JSON parsing issues + except ImportError as e: + missing_deps.append(f"httpx: {e}") + + try: + import mcp + # Silent check - no print to avoid JSON parsing issues + except ImportError as e: + missing_deps.append(f"mcp: {e}") + + try: + from PIL import Image + # Silent check - no print to avoid JSON parsing issues + except ImportError as e: + missing_deps.append(f"Pillow: {e}") + + try: + import dotenv + # Silent check - no print to avoid JSON parsing issues + except ImportError as e: + missing_deps.append(f"python-dotenv: {e}") + + try: + import pydantic + # Silent check - no print to avoid JSON parsing issues + except ImportError as e: + missing_deps.append(f"pydantic: {e}") + + if missing_deps: + # Log to file instead of print to avoid JSON parsing issues + logger = logging.getLogger(__name__) + logger.error(f"Missing dependencies: {missing_deps}") + return False + + return True + + +def check_local_imports(): + """Check if local modules can be imported - silently check for MCP compatibility""" + logger = logging.getLogger(__name__) + + try: + from src.connector import Config + # Silent check - no print to avoid JSON parsing issues + except ImportError as e: + logger.error(f"Failed to import local Config: {e}") + return False + + try: + from src.server import main + # Silent check - no print to avoid JSON parsing issues + return True + except ImportError as e: + logger.error(f"Failed to import local server main: {e}", exc_info=True) + return False + + +if __name__ == "__main__": + setup_logging() + logger = logging.getLogger(__name__) + + # Check dependencies first - silently for MCP compatibility + if not check_dependencies(): + logger.error("Dependencies check failed") + sys.exit(1) + + # Check local imports - silently for MCP compatibility + if not check_local_imports(): + logger.error("Local imports check failed") + sys.exit(1) + + # Start the MCP server without console output to avoid JSON parsing issues + try: + # Import main function after checks + from src.server import main + logger.info("Starting FLUX.1 Edit MCP Server...") + asyncio.run(main()) + except KeyboardInterrupt: + logger.info("Server stopped by user") + sys.exit(0) + except Exception as e: + logger.error(f"Server error: {e}", exc_info=True) + sys.exit(1) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..cc834df --- /dev/null +++ b/requirements.txt @@ -0,0 +1,29 @@ +# FLUX.1 Edit MCP Server Dependencies + +# Core MCP Server +mcp==1.1.0 + +# HTTP Client for FLUX API +httpx==0.28.1 +aiohttp==3.11.7 + +# Image Processing +Pillow==11.0.0 + +# Environment and Configuration +python-dotenv==1.0.1 + +# Data Validation +pydantic==2.10.3 + +# Async utilities - asyncio is built into Python 3.7+ +# asyncio-compat package not needed for modern Python versions + +# Logging +structlog==24.4.0 + +# Development and Testing (optional) +pytest==8.3.4 +pytest-asyncio==0.25.0 +pytest-mock==3.14.0 +black==24.10.0 diff --git a/run.bat b/run.bat new file mode 100644 index 0000000..2f4f831 --- /dev/null +++ b/run.bat @@ -0,0 +1,74 @@ +@echo off +echo Starting FLUX.1 Edit MCP Server... +echo ================================ + +REM Check if Python is available +python --version >nul 2>&1 +if errorlevel 1 ( + echo Error: Python is not installed or not in PATH + echo Please install Python 3.8 or higher + pause + exit /b 1 +) + +REM Check if virtual environment exists +if not exist "venv" ( + echo Creating virtual environment... + python -m venv venv + if errorlevel 1 ( + echo Error: Failed to create virtual environment + pause + exit /b 1 + ) +) + +REM Activate virtual environment +echo Activating virtual environment... +call venv\Scripts\activate.bat +if errorlevel 1 ( + echo Error: Failed to activate virtual environment + pause + exit /b 1 +) + +REM Install/upgrade dependencies +echo Installing dependencies... +pip install -r requirements.txt +if errorlevel 1 ( + echo Error: Failed to install dependencies + pause + exit /b 1 +) + +REM Check if .env file exists +if not exist ".env" ( + echo Warning: .env file not found + echo Please copy .env.example to .env and configure your FLUX_API_KEY + echo. + echo Creating .env from example... + copy .env.example .env + echo. + echo Please edit .env file and add your FLUX_API_KEY before running the server + pause + exit /b 1 +) + +REM Create directories if they don't exist +if not exist "input_images" mkdir input_images +if not exist "generated_images" mkdir generated_images +if not exist "temp" mkdir temp + +echo. +echo Starting FLUX.1 Edit MCP Server... +echo Press Ctrl+C to stop the server +echo. + +REM Run the server with UTF-8 encoding to prevent Unicode errors +chcp 65001 >nul 2>&1 +set PYTHONIOENCODING=utf-8 +set PYTHONUTF8=1 +python main.py + +echo. +echo Server stopped. +pause diff --git a/run.ps1 b/run.ps1 new file mode 100644 index 0000000..a7fac9b --- /dev/null +++ b/run.ps1 @@ -0,0 +1,83 @@ +# PowerShell script to run FLUX.1 Edit MCP Server + +Write-Host "Starting FLUX.1 Edit MCP Server..." -ForegroundColor Green +Write-Host "================================" -ForegroundColor Green + +# Check if Python is available +try { + $pythonVersion = python --version + Write-Host "Found Python: $pythonVersion" -ForegroundColor Blue +} catch { + Write-Host "Error: Python is not installed or not in PATH" -ForegroundColor Red + Write-Host "Please install Python 3.8 or higher" -ForegroundColor Yellow + Read-Host "Press Enter to exit" + exit 1 +} + +# Check if virtual environment exists +if (-not (Test-Path "venv")) { + Write-Host "Creating virtual environment..." -ForegroundColor Yellow + python -m venv venv + if ($LASTEXITCODE -ne 0) { + Write-Host "Error: Failed to create virtual environment" -ForegroundColor Red + Read-Host "Press Enter to exit" + exit 1 + } +} + +# Activate virtual environment +Write-Host "Activating virtual environment..." -ForegroundColor Blue +& "venv\Scripts\Activate.ps1" +if ($LASTEXITCODE -ne 0) { + Write-Host "Error: Failed to activate virtual environment" -ForegroundColor Red + Write-Host "You might need to enable PowerShell script execution:" -ForegroundColor Yellow + Write-Host "Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser" -ForegroundColor Yellow + Read-Host "Press Enter to exit" + exit 1 +} + +# Install/upgrade dependencies +Write-Host "Installing dependencies..." -ForegroundColor Blue +pip install -r requirements.txt +if ($LASTEXITCODE -ne 0) { + Write-Host "Error: Failed to install dependencies" -ForegroundColor Red + Read-Host "Press Enter to exit" + exit 1 +} + +# Check if .env file exists +if (-not (Test-Path ".env")) { + Write-Host "Warning: .env file not found" -ForegroundColor Yellow + Write-Host "Please copy .env.example to .env and configure your FLUX_API_KEY" -ForegroundColor Yellow + Write-Host "" + Write-Host "Creating .env from example..." -ForegroundColor Blue + Copy-Item ".env.example" ".env" + Write-Host "" + Write-Host "Please edit .env file and add your FLUX_API_KEY before running the server" -ForegroundColor Red + Read-Host "Press Enter to exit" + exit 1 +} + +# Create directories if they don't exist +@("input_images", "generated_images", "temp") | ForEach-Object { + if (-not (Test-Path $_)) { + New-Item -ItemType Directory -Path $_ | Out-Null + Write-Host "Created directory: $_" -ForegroundColor Blue + } +} + +Write-Host "" +Write-Host "Starting FLUX.1 Edit MCP Server..." -ForegroundColor Green +Write-Host "Press Ctrl+C to stop the server" -ForegroundColor Yellow +Write-Host "" + +# Run the server +try { + python main.py +} catch { + Write-Host "Error running server: $_" -ForegroundColor Red +} + +Write-Host "" +Write-Host "Server stopped." -ForegroundColor Yellow +Read-Host "Press Enter to exit" diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..163143c --- /dev/null +++ b/src/__init__.py @@ -0,0 +1,5 @@ +"""FLUX.1 Edit MCP Server package""" + +__version__ = "1.0.0" +__author__ = "FLUX.1 Edit Team" +__description__ = "MCP Server for FLUX.1 Kontext image editing" diff --git a/src/connector/__init__.py b/src/connector/__init__.py new file mode 100644 index 0000000..5ac8209 --- /dev/null +++ b/src/connector/__init__.py @@ -0,0 +1,33 @@ +"""Connector package for FLUX.1 Edit""" + +import sys +import logging + +logger = logging.getLogger(__name__) + +try: + from .config import Config + logger.debug("Config imported successfully") +except ImportError as e: + logger.error(f"Failed to import Config: {e}") + raise + +try: + from .flux_client import FluxEditClient, FluxEditRequest, FluxEditResponse + logger.debug("FluxEditClient classes imported successfully") +except ImportError as e: + logger.error(f"Failed to import FluxEditClient classes: {e}") + # Check if the issue is with aiohttp specifically + try: + import aiohttp + except ImportError: + logger.error("aiohttp is not installed. Please run: pip install aiohttp==3.11.7") + print("\n❌ Missing dependency: aiohttp") + print("Please run one of the following commands:") + print(" - install_dependencies.bat (on Windows)") + print(" - pip install -r requirements.txt") + print(" - pip install aiohttp==3.11.7") + sys.exit(1) + raise + +__all__ = ['Config', 'FluxEditClient', 'FluxEditRequest', 'FluxEditResponse'] diff --git a/src/connector/config.py b/src/connector/config.py new file mode 100644 index 0000000..1ee400a --- /dev/null +++ b/src/connector/config.py @@ -0,0 +1,272 @@ +"""Configuration module for FLUX.1 Edit""" + +import os +import logging +import random +from typing import Optional, Tuple +from pathlib import Path +from datetime import datetime +from dotenv import load_dotenv + +logger = logging.getLogger(__name__) + + +class Config: + """Configuration class for FLUX.1 Edit""" + + # FLUX.1 Kontext API Configuration + API_BASE_URL = "https://api.bfl.ai" + EDIT_ENDPOINT = "/flux-kontext-pro" + RESULT_ENDPOINT = "/v1/get_result" + + # Fixed FLUX parameters based on requirements + MODEL_NAME = "flux-kontext-pro" + OUTPUT_FORMAT = "png" + PROMPT_UPSAMPLING = False + DEFAULT_SAFETY_TOLERANCE = 2 + + # Image size limits + MAX_IMAGE_SIZE_MB = 20 # FLUX.1 Kontext limit + + # Aspect ratios supported + SUPPORTED_ASPECT_RATIOS = [ + "1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21" + ] + DEFAULT_ASPECT_RATIO = "1:1" + + def __init__(self): + """Initialize configuration""" + # Load environment variables + env_path = Path(__file__).parent.parent.parent / '.env' + if env_path.exists(): + load_dotenv(env_path) + logger.info(f"Loaded environment from {env_path}") + else: + logger.warning(f"No .env file found at {env_path}") + + # API Configuration + self.api_key = os.getenv('FLUX_API_KEY', '') + self.api_base_url = os.getenv('FLUX_API_BASE_URL', self.API_BASE_URL) + + # Server Configuration + self.log_level = os.getenv('LOG_LEVEL', 'INFO') + self.max_image_size_mb = int(os.getenv('MAX_IMAGE_SIZE_MB', str(self.MAX_IMAGE_SIZE_MB))) + self.default_timeout = int(os.getenv('DEFAULT_TIMEOUT', '300')) + self.polling_interval = int(os.getenv('POLLING_INTERVAL', '2')) + self.max_polling_attempts = int(os.getenv('MAX_POLLING_ATTEMPTS', '150')) + + # Default Settings + self.default_aspect_ratio = os.getenv('DEFAULT_ASPECT_RATIO', self.DEFAULT_ASPECT_RATIO) + self.safety_tolerance = int(os.getenv('DEFAULT_SAFETY_TOLERANCE', str(self.DEFAULT_SAFETY_TOLERANCE))) + self.prompt_upsampling = os.getenv('PROMPT_UPSAMPLING', 'false').lower() == 'true' + + # Paths Configuration + self.base_path = Path(__file__).parent.parent.parent + + # Input directory for reading source images + default_input_dir = str(self.base_path / 'input_images') + self.input_path = Path(os.getenv('INPUT_PATH', default_input_dir)) + + # Single output directory for everything + default_output_dir = str(self.base_path / 'generated_images') + self.generated_images_path = Path(os.getenv('GENERATED_IMAGES_PATH', default_output_dir)) + + # File naming configuration + self.output_filename_prefix = os.getenv('OUTPUT_FILENAME_PREFIX', 'fluxedit') + self.save_parameters = os.getenv('SAVE_PARAMETERS', 'true').lower() == 'true' + + # Ensure all required directories exist with proper error handling + self._ensure_directories() + + logger.info(f"Input path: {self.input_path}") + logger.info(f"Generated images path: {self.generated_images_path}") + + def _ensure_directories(self) -> None: + """ + Ensure all required directories exist with proper permissions and error handling + """ + directories = [ + ("input_images", self.input_path), + ("generated_images", self.generated_images_path) + ] + + for dir_name, dir_path in directories: + try: + # Create directory with parents if needed + dir_path.mkdir(parents=True, exist_ok=True) + + # Verify directory is accessible + if not dir_path.exists(): + raise RuntimeError(f"Failed to create {dir_name} directory: {dir_path}") + + if not dir_path.is_dir(): + raise RuntimeError(f"{dir_name} path exists but is not a directory: {dir_path}") + + # Test write permissions by creating a temporary test file + test_file = dir_path / ".fluxedit_test_write" + try: + test_file.touch() + test_file.unlink() # Delete test file + except PermissionError: + raise RuntimeError(f"No write permission for {dir_name} directory: {dir_path}") + except Exception as e: + logger.warning(f"Could not test write permissions for {dir_name} directory: {e}") + + logger.debug(f"✅ {dir_name.title()} directory ready: {dir_path}") + + except Exception as e: + logger.error(f"❌ Failed to setup {dir_name} directory ({dir_path}): {e}") + raise RuntimeError(f"Directory setup failed for {dir_name}: {e}") from e + + def ensure_output_directory(self) -> None: + """ + Runtime method to ensure output directory exists (in case it gets deleted) + """ + try: + if not self.generated_images_path.exists(): + logger.warning(f"Output directory missing, recreating: {self.generated_images_path}") + self.generated_images_path.mkdir(parents=True, exist_ok=True) + + # Verify creation was successful + if not self.generated_images_path.exists(): + raise RuntimeError(f"Failed to recreate output directory: {self.generated_images_path}") + + logger.info(f"✅ Output directory recreated: {self.generated_images_path}") + + except Exception as e: + logger.error(f"❌ Failed to ensure output directory: {e}") + raise + + def generate_base_name_simple(self) -> str: + """ + Generate simple base name in format: fluxedit_{yyyymmdd}_{hhmmss} + + Returns: + str: Base name for files + """ + now = datetime.now() + date_str = now.strftime("%Y%m%d") + time_str = now.strftime("%H%M%S") + + return f"{self.output_filename_prefix}_{date_str}_{time_str}" + + def generate_base_name(self, seed: Optional[int] = None) -> str: + """ + Generate base name with seed for files + + Args: + seed: Optional seed value, generated if None + + Returns: + str: Base name for files + """ + if seed is None: + seed = random.randint(0, 999999) + + now = datetime.now() + date_str = now.strftime("%Y%m%d") + time_str = now.strftime("%H%M%S") + + return f"{self.output_filename_prefix}_{seed}_{date_str}_{time_str}" + + def generate_filename(self, base_name: str, file_number: int = 1, extension: str = 'png') -> str: + """ + Generate filename from base name: + - Input: fluxedit_{seed}_{yyyymmdd}_{hhmmss}_000.{ext} + - Output: fluxedit_{seed}_{yyyymmdd}_{hhmmss}_001.png + - JSON: fluxedit_{seed}_{yyyymmdd}_{hhmmss}_001.json + + Args: + base_name: Base name (e.g., fluxedit_123456_20250826_143022) + file_number: File number (0 for input, 1+ for outputs) + extension: File extension + + Returns: + str: Generated filename + """ + return f"{base_name}_{file_number:03d}.{extension}" + + def get_output_path(self, base_name: str, file_number: int = 1, extension: str = 'png') -> Path: + """ + Get full path for output file with directory verification + + Args: + base_name: Base name for the file + file_number: File number (0 for input, 1+ for outputs) + extension: File extension + + Returns: + Path: Full path to the file + """ + # Ensure output directory exists before returning path + self.ensure_output_directory() + + filename = self.generate_filename(base_name, file_number, extension) + return self.generated_images_path / filename + + def get_api_url(self, endpoint: str) -> str: + """ + Get full API URL for endpoint + + Args: + endpoint: API endpoint (e.g., '/flux-kontext-pro') + + Returns: + str: Full API URL + """ + return f"{self.api_base_url}{endpoint}" + + def validate(self) -> bool: + """ + Validate configuration + + Returns: + bool: True if configuration is valid + """ + if not self.api_key: + logger.error("FLUX_API_KEY is not set") + return False + + if self.max_image_size_mb <= 0 or self.max_image_size_mb > 50: + logger.error(f"Invalid MAX_IMAGE_SIZE_MB: {self.max_image_size_mb} (must be 1-50)") + return False + + if self.default_timeout <= 0: + logger.error(f"Invalid DEFAULT_TIMEOUT: {self.default_timeout}") + return False + + if self.polling_interval <= 0: + logger.error(f"Invalid POLLING_INTERVAL: {self.polling_interval}") + return False + + if self.max_polling_attempts <= 0: + logger.error(f"Invalid MAX_POLLING_ATTEMPTS: {self.max_polling_attempts}") + return False + + if self.default_aspect_ratio not in self.SUPPORTED_ASPECT_RATIOS: + logger.error(f"Invalid DEFAULT_ASPECT_RATIO: {self.default_aspect_ratio}") + return False + + logger.info("Configuration validated successfully") + return True + + def get_max_image_size_bytes(self) -> int: + """Get maximum image size in bytes""" + return self.max_image_size_mb * 1024 * 1024 + + def __str__(self) -> str: + """String representation""" + return ( + f"FLUX.1 Edit Configuration:\n" + f" API Key: {'***' + self.api_key[-4:] if self.api_key else 'Not Set'}\n" + f" API Base URL: {self.api_base_url}\n" + f" Max Image Size: {self.max_image_size_mb}MB\n" + f" Timeout: {self.default_timeout}s\n" + f" Polling Interval: {self.polling_interval}s\n" + f" Max Polling Attempts: {self.max_polling_attempts}\n" + f" Default Aspect Ratio: {self.default_aspect_ratio}\n" + f" Safety Tolerance: {self.safety_tolerance}\n" + f" Input Directory: {self.input_path}\n" + f" Output Directory: {self.generated_images_path}\n" + f" Save Parameters: {self.save_parameters}" + ) diff --git a/src/connector/flux_client.py b/src/connector/flux_client.py new file mode 100644 index 0000000..8504ce6 --- /dev/null +++ b/src/connector/flux_client.py @@ -0,0 +1,334 @@ +"""FLUX.1 Kontext API Client for image editing""" + +import json +import logging +import asyncio +import base64 +from typing import Optional, Dict, Any, Tuple +from dataclasses import dataclass +from datetime import datetime +import aiohttp + +from .config import Config + +logger = logging.getLogger(__name__) + + +@dataclass +class FluxEditRequest: + """Request data for FLUX.1 Kontext image editing""" + 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 + + +@dataclass +class FluxEditResponse: + """Response data from FLUX.1 Kontext image editing""" + success: bool + edited_image_data: Optional[bytes] = None + image_size: Optional[Tuple[int, int]] = None + execution_time: float = 0.0 + request_id: Optional[str] = None + result_url: Optional[str] = None + error_message: Optional[str] = None + metadata: Optional[Dict[str, Any]] = None + + +class FluxEditClient: + """FLUX.1 Kontext API client for image editing""" + + def __init__(self, config: Config): + """Initialize client with configuration""" + self.config = config + self.session = None + + async def __aenter__(self): + """Async context manager entry""" + await self._ensure_session() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Async context manager exit""" + await self.close() + + 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) + self.session = aiohttp.ClientSession(timeout=timeout) + + async def close(self): + """Close the HTTP session""" + if self.session and not self.session.closed: + await self.session.close() + self.session = None + + def _get_headers(self) -> Dict[str, str]: + """Get request headers with API key""" + return { + 'Content-Type': 'application/json', + 'X-Key': self.config.api_key + } + + async def _create_edit_request(self, request: FluxEditRequest) -> Optional[str]: + """ + Create image edit request and return request_id + + Args: + request: FluxEditRequest with edit parameters + + Returns: + str: request_id if successful, None otherwise + """ + try: + await self._ensure_session() + + # Prepare request payload + payload = { + "prompt": request.prompt, + "input_image": request.input_image_b64, + "seed": request.seed, + "safety_tolerance": request.safety_tolerance, + "output_format": request.output_format + } + + # 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 + + 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())}") + + 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: + logger.info(f"Edit request created successfully: {request_id}") + return request_id + else: + logger.error(f"No request_id in response: {result}") + return None + else: + error_text = await response.text() + logger.error(f"Failed to create edit request: {response.status} - {error_text}") + return None + + except asyncio.TimeoutError: + logger.error("Timeout creating edit request") + 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]]: + """ + Poll for edit result using request_id + + Args: + request_id: Request ID from create_edit_request + + Returns: + dict: Result data if successful, None otherwise + """ + try: + await self._ensure_session() + + url = self.config.get_api_url(self.config.RESULT_ENDPOINT) + params = {"id": request_id} + + attempts = 0 + max_attempts = self.config.max_polling_attempts + interval = self.config.polling_interval + + logger.info(f"Starting to poll for result: {request_id}") + + while attempts < max_attempts: + try: + async with self.session.get(url, params=params, headers=self._get_headers()) as response: + if response.status == 200: + result = await response.json() + + # Check result status + status = result.get('status', '').lower() + + if status == 'ready': + logger.info(f"Result ready after {attempts + 1} attempts") + return result + elif status in ['failed', 'error']: + error_msg = result.get('error', 'Unknown error') + logger.error(f"Edit failed: {error_msg}") + return None + else: + # Still processing, continue polling + logger.debug(f"Status: {status}, continuing to poll...") + else: + logger.warning(f"Polling attempt {attempts + 1} failed: {response.status}") + + attempts += 1 + if attempts < max_attempts: + await asyncio.sleep(interval) + + except asyncio.TimeoutError: + logger.warning(f"Timeout on polling attempt {attempts + 1}") + attempts += 1 + if attempts < max_attempts: + await asyncio.sleep(interval) + except Exception as e: + logger.error(f"Error during polling attempt {attempts + 1}: {e}") + attempts += 1 + if attempts < max_attempts: + await asyncio.sleep(interval) + + logger.error(f"Polling timeout after {max_attempts} attempts") + return None + + except Exception as e: + logger.error(f"Error polling for result: {e}", exc_info=True) + return None + + async def _download_result_image(self, result_url: str) -> Optional[bytes]: + """ + Download result image from signed URL + + Args: + result_url: Signed URL for the result image + + Returns: + bytes: Image data if successful, None otherwise + """ + try: + await self._ensure_session() + + logger.info(f"Downloading result image from URL") + + async with self.session.get(result_url) as response: + if response.status == 200: + image_data = await response.read() + logger.info(f"Downloaded image: {len(image_data)} bytes") + return image_data + else: + logger.error(f"Failed to download image: {response.status}") + return None + + except Exception as e: + logger.error(f"Error downloading result image: {e}", exc_info=True) + return None + + def _get_image_size(self, image_data: bytes) -> Optional[Tuple[int, int]]: + """ + Get image dimensions from image data + + Args: + image_data: Image bytes + + Returns: + tuple: (width, height) if successful, None otherwise + """ + try: + from PIL import Image + import io + + with Image.open(io.BytesIO(image_data)) as img: + return img.size + except Exception as e: + logger.warning(f"Could not determine image size: {e}") + return None + + async def edit_image(self, request: FluxEditRequest) -> FluxEditResponse: + """ + Edit image using FLUX.1 Kontext API + + Args: + request: FluxEditRequest with all parameters + + Returns: + FluxEditResponse: Response with edited image or error + """ + start_time = datetime.now() + + try: + 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: + return FluxEditResponse( + success=False, + error_message="Failed to create edit request", + execution_time=(datetime.now() - start_time).total_seconds() + ) + + # Step 2: Poll for result + result = await self._poll_result(request_id) + if not result: + return FluxEditResponse( + success=False, + request_id=request_id, + error_message="Failed to get edit result or polling timeout", + execution_time=(datetime.now() - start_time).total_seconds() + ) + + # Step 3: Download result image + result_url = result.get('result', {}).get('sample') + if not result_url: + return FluxEditResponse( + success=False, + request_id=request_id, + error_message="No result URL in response", + execution_time=(datetime.now() - start_time).total_seconds() + ) + + image_data = await self._download_result_image(result_url) + if not image_data: + return FluxEditResponse( + success=False, + request_id=request_id, + result_url=result_url, + error_message="Failed to download result image", + execution_time=(datetime.now() - start_time).total_seconds() + ) + + # Get image dimensions + image_size = self._get_image_size(image_data) + + execution_time = (datetime.now() - start_time).total_seconds() + + logger.info(f"FLUX edit completed successfully in {execution_time:.1f}s") + + return FluxEditResponse( + success=True, + edited_image_data=image_data, + image_size=image_size, + execution_time=execution_time, + request_id=request_id, + result_url=result_url, + metadata={ + "seed": request.seed, + "aspect_ratio": request.aspect_ratio, + "safety_tolerance": request.safety_tolerance, + "prompt_upsampling": request.prompt_upsampling + } + ) + + except Exception as e: + execution_time = (datetime.now() - start_time).total_seconds() + logger.error(f"FLUX edit failed: {e}", exc_info=True) + + return FluxEditResponse( + success=False, + error_message=f"Unexpected error: {str(e)}", + execution_time=execution_time + ) + finally: + # Ensure session cleanup (optional, can be managed by context manager) + pass diff --git a/src/server/__init__.py b/src/server/__init__.py new file mode 100644 index 0000000..ea6bd16 --- /dev/null +++ b/src/server/__init__.py @@ -0,0 +1,7 @@ +"""Server package for FLUX.1 Edit""" + +from .mcp_server import FluxEditMCPServer, create_server, main +from .handlers import ToolHandlers +from .models import TOOL_DEFINITIONS, ToolName + +__all__ = ['FluxEditMCPServer', 'create_server', 'main', 'ToolHandlers', 'TOOL_DEFINITIONS', 'ToolName'] diff --git a/src/server/handlers.py b/src/server/handlers.py new file mode 100644 index 0000000..445560e --- /dev/null +++ b/src/server/handlers.py @@ -0,0 +1,605 @@ +"""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"[ERROR] 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"[ERROR] 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"[SUCCESS] 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"\nOutput: {Path(saved_path).name}" + text += f"\nInput: {Path(input_generated_path).name}" + if json_path: + text += f"\nParameters: {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"[ERROR] 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"[ERROR] 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"[ERROR] 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 += "WARNING: 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"[ERROR] 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"[ERROR] 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"[ERROR] 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"[ERROR] 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"[SUCCESS] 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"\nOutput: {Path(saved_path).name}" + text += f"\nInput copy: {Path(input_generated_path).name}" + if json_path: + text += f"\nParameters: {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"[ERROR] 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"[ERROR] 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"[SUCCESS] 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"[ERROR] 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"[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/handlers_backup.py b/src/server/handlers_backup.py new file mode 100644 index 0000000..87ab4bd --- /dev/null +++ b/src/server/handlers_backup.py @@ -0,0 +1,605 @@ +"""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)}" + )] diff --git a/src/server/mcp_server.py b/src/server/mcp_server.py new file mode 100644 index 0000000..1a1566c --- /dev/null +++ b/src/server/mcp_server.py @@ -0,0 +1,178 @@ +"""MCP Server for FLUX.1 Edit""" + +import logging +import asyncio +from typing import Dict, Any, List + +import mcp.types as types +from mcp import server +from mcp.server import Server +from mcp.server.models import InitializationOptions + +from ..connector import Config +from .models import TOOL_DEFINITIONS, ToolName +from .handlers import ToolHandlers + +logger = logging.getLogger(__name__) + + +class FluxEditMCPServer: + """FLUX.1 Edit MCP Server""" + + def __init__(self): + """Initialize the MCP server""" + self.config = Config() + self.app = Server("flux1-edit") + self.handlers = ToolHandlers(self.config) + self._setup_handlers() + + logger.info("FLUX.1 Edit MCP Server initialized") + + def _setup_handlers(self): + """Setup MCP server handlers""" + + # List tools handler + @self.app.list_tools() + async def list_tools() -> List[types.Tool]: + """List available tools""" + tools = [] + + for tool_name, tool_def in TOOL_DEFINITIONS.items(): + # Build properties for parameters + properties = {} + required = [] + + for param in tool_def.parameters: + prop_def = { + "type": param.type, + "description": param.description + } + + # Add enum if specified + if param.enum: + prop_def["enum"] = param.enum + + # Add default if specified + if param.default is not None: + prop_def["default"] = param.default + + properties[param.name] = prop_def + + if param.required: + required.append(param.name) + + # Build tool schema + tool = types.Tool( + name=tool_def.name, + description=tool_def.description, + inputSchema={ + "type": "object", + "properties": properties, + "required": required + } + ) + tools.append(tool) + + logger.debug(f"Listed {len(tools)} tools") + return tools + + # Call tool handler + @self.app.call_tool() + async def call_tool(name: str, arguments: Dict[str, Any]) -> List[types.TextContent | types.ImageContent]: + """Handle tool calls""" + try: + logger.info(f"Tool call: {name}") + logger.debug(f"Arguments: {self._sanitize_args_for_logging(arguments)}") + + # Route to appropriate handler + if name == ToolName.FLUX_EDIT_IMAGE: + return await self.handlers.handle_flux_edit_image(arguments) + elif name == ToolName.FLUX_EDIT_IMAGE_FROM_FILE: + 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: + return [types.TextContent( + type="text", + text=f"[ERROR] Unknown tool: {name}" + )] + + except Exception as e: + logger.error(f"Error calling tool {name}: {e}", exc_info=True) + return [types.TextContent( + type="text", + text=f"[ERROR] Tool execution error: {str(e)}" + )] + + def _sanitize_args_for_logging(self, args: Dict[str, Any]) -> Dict[str, Any]: + """ + Remove sensitive data from arguments for logging + + Args: + args: Original arguments + + Returns: + dict: Sanitized arguments + """ + safe_args = args.copy() + + # Don't log full base64 image data + if 'input_image_b64' in safe_args: + b64_data = safe_args['input_image_b64'] + safe_args['input_image_b64'] = f"" + + # Truncate long prompts + if 'prompt' in safe_args and len(safe_args['prompt']) > 100: + safe_args['prompt'] = safe_args['prompt'][:100] + '...' + + return safe_args + + def validate_config(self) -> bool: + """Validate server configuration""" + return self.config.validate() + + async def run(self): + """Run the MCP server""" + if not self.validate_config(): + logger.error("Configuration validation failed") + return False + + logger.info("Starting FLUX.1 Edit MCP Server...") + logger.info(f"Configuration:\n{self.config}") + + # Run the server + async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): + await self.app.run( + read_stream, + write_stream, + InitializationOptions( + server_name="flux1-edit", + server_version="1.0.0" + ) + ) + + +# Server factory function +def create_server() -> FluxEditMCPServer: + """Create and return a FLUX.1 Edit MCP Server instance""" + return FluxEditMCPServer() + + +# Main function for running the server +async def main(): + """Main entry point for the server""" + # Setup logging + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + ) + + # Create and run server + server_instance = create_server() + await server_instance.run() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/src/server/models.py b/src/server/models.py new file mode 100644 index 0000000..f298e4e --- /dev/null +++ b/src/server/models.py @@ -0,0 +1,190 @@ +"""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 diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..1caca02 --- /dev/null +++ b/src/utils/__init__.py @@ -0,0 +1,55 @@ +"""Utils package for FLUX.1 Edit""" + +from .image_utils import ( + get_file_size_mb, + get_image_size_from_bytes, + validate_image_file, + get_image_dimensions, + get_image_dimensions_from_bytes, + save_image, + encode_image_base64, + decode_image_base64, + optimize_image_for_flux, + convert_image_to_base64, + validate_aspect_ratio, + get_optimal_aspect_ratio +) + +from .validation import ( + validate_edit_parameters, + validate_file_parameters, + validate_move_file_parameters, + validate_image_path_parameter, + sanitize_prompt, + validate_aspect_ratio_format, + validate_seed_range, + validate_base64_image, + validate_filename_safety +) + +__all__ = [ + # Image utilities + 'get_file_size_mb', + 'get_image_size_from_bytes', + 'validate_image_file', + 'get_image_dimensions', + 'get_image_dimensions_from_bytes', + 'save_image', + 'encode_image_base64', + 'decode_image_base64', + 'optimize_image_for_flux', + 'convert_image_to_base64', + 'validate_aspect_ratio', + 'get_optimal_aspect_ratio', + + # Validation utilities + 'validate_edit_parameters', + 'validate_file_parameters', + 'validate_move_file_parameters', + 'validate_image_path_parameter', + 'sanitize_prompt', + 'validate_aspect_ratio_format', + 'validate_seed_range', + 'validate_base64_image', + 'validate_filename_safety' +] diff --git a/src/utils/image_utils.py b/src/utils/image_utils.py new file mode 100644 index 0000000..d198a54 --- /dev/null +++ b/src/utils/image_utils.py @@ -0,0 +1,385 @@ +"""Image utility functions for FLUX.1 Edit MCP Server""" + +import base64 +import io +import logging +from pathlib import Path +from typing import Tuple, Optional, Union + +try: + from PIL import Image +except ImportError: + raise ImportError("Pillow is required. Install with: pip install pillow") + +logger = logging.getLogger(__name__) + + +def get_file_size_mb(file_path: Union[str, Path]) -> float: + """Get file size in MB""" + path = Path(file_path) + if path.exists(): + return path.stat().st_size / (1024 * 1024) + return 0.0 + + +def get_image_size_from_bytes(data: bytes) -> float: + """Get size of image data in MB""" + return len(data) / (1024 * 1024) + + +def validate_image_file(file_path: str, max_size_mb: int = 20) -> Tuple[bool, float, Optional[str]]: + """ + Validate an image file for FLUX.1 Kontext (20MB limit) + + Args: + file_path: Path to image file + max_size_mb: Maximum file size in MB (default: 20 for FLUX) + + Returns: + tuple: (is_valid, size_mb, error_message) + """ + try: + path = Path(file_path) + + # Check if file exists + if not path.exists(): + return False, 0, f"File not found: {file_path}" + + # Check file size + size_mb = get_file_size_mb(path) + + # Check if it's a valid image + try: + with Image.open(file_path) as img: + # Verify it's a supported format + if img.format not in ['PNG', 'JPEG', 'JPG', 'GIF', 'BMP', 'WEBP', 'TIFF']: + return False, size_mb, f"Unsupported format: {img.format}" + + # Check dimensions (reasonable limits) + width, height = img.size + if width > 8192 or height > 8192: + return False, size_mb, f"Image dimensions too large: {width}x{height} (max: 8192x8192)" + + if width < 32 or height < 32: + return False, size_mb, f"Image dimensions too small: {width}x{height} (min: 32x32)" + + except Exception as e: + return False, size_mb, f"Invalid image file: {str(e)}" + + # Check size limit + if size_mb > max_size_mb: + return False, size_mb, f"File size {size_mb:.2f}MB exceeds {max_size_mb}MB limit" + + return True, size_mb, None + + except Exception as e: + logger.error(f"Error validating image: {e}") + return False, 0, str(e) + + +def get_image_dimensions(file_path: str) -> Tuple[int, int]: + """ + Get image dimensions + + Args: + file_path: Path to image file + + Returns: + tuple: (width, height) + """ + try: + with Image.open(file_path) as img: + return img.size + except Exception as e: + logger.error(f"Error getting dimensions: {e}") + return (0, 0) + + +def get_image_dimensions_from_bytes(image_data: bytes) -> Tuple[int, int]: + """ + Get image dimensions from bytes + + Args: + image_data: Image data as bytes + + Returns: + tuple: (width, height) + """ + try: + with Image.open(io.BytesIO(image_data)) as img: + return img.size + except Exception as e: + logger.error(f"Error getting dimensions from bytes: {e}") + return (0, 0) + + +def save_image(image_data: bytes, output_path: str) -> bool: + """ + Save image data to file + + Args: + image_data: Image data as bytes + output_path: Output file path + + Returns: + bool: Success status + """ + try: + path = Path(output_path) + + # Ensure parent directory exists + try: + path.parent.mkdir(parents=True, exist_ok=True) + + # Verify directory was created successfully + if not path.parent.exists(): + logger.error(f"Failed to create directory: {path.parent}") + return False + + # Check if parent is actually a directory + if not path.parent.is_dir(): + logger.error(f"Parent path exists but is not a directory: {path.parent}") + return False + + except PermissionError as e: + logger.error(f"Permission denied creating directory {path.parent}: {e}") + return False + except Exception as e: + logger.error(f"Failed to create directory {path.parent}: {e}") + return False + + # Save the image file + try: + with open(path, 'wb') as f: + f.write(image_data) + + # Verify file was written successfully + if not path.exists() or path.stat().st_size != len(image_data): + logger.error(f"File save verification failed: {path}") + return False + + logger.info(f"Image saved: {path} ({len(image_data):,} bytes)") + return True + + except PermissionError as e: + logger.error(f"Permission denied writing file {path}: {e}") + return False + except OSError as e: + logger.error(f"OS error writing file {path}: {e}") + return False + + except Exception as e: + logger.error(f"Unexpected error saving image to {output_path}: {e}") + return False + + +def encode_image_base64(image_data: bytes) -> str: + """ + Encode image data to base64 string + + Args: + image_data: Image bytes + + Returns: + str: Base64 encoded string + """ + return base64.b64encode(image_data).decode('utf-8') + + +def decode_image_base64(base64_str: str) -> bytes: + """ + Decode base64 string to image data + Supports both raw base64 and data URL formats + + Args: + base64_str: Base64 encoded string (with or without data URL prefix) + + Returns: + bytes: Image data + """ + # Handle data URL format (e.g., "data:image/jpeg;base64,...") + if base64_str.startswith('data:'): + # Find the comma that separates the header from data + comma_index = base64_str.find(',') + if comma_index != -1: + base64_str = base64_str[comma_index + 1:] + else: + raise ValueError("Invalid data URL format: no comma found") + + # Remove any whitespace/newlines + base64_str = base64_str.strip().replace('\n', '').replace('\r', '') + + try: + return base64.b64decode(base64_str) + except Exception as e: + raise ValueError(f"Failed to decode base64 data: {e}") + + +def optimize_image_for_flux(image_path: str, max_size_mb: float = 20.0) -> bytes: + """ + Optimize image for FLUX.1 Kontext API (20MB limit) + + Args: + image_path: Path to input image + max_size_mb: Maximum size in MB (default: 20 for FLUX) + + Returns: + bytes: Optimized image data + """ + max_size_bytes = max_size_mb * 1024 * 1024 + + try: + with Image.open(image_path) as img: + # For FLUX, we want to preserve quality as much as possible + # since 20MB is quite generous + + # Convert to RGB if needed (FLUX typically prefers RGB) + if img.mode != 'RGB': + if img.mode == 'RGBA': + # Create white background for transparent images + background = Image.new('RGB', img.size, (255, 255, 255)) + background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None) + img = background + else: + img = img.convert('RGB') + + # Try PNG first (lossless) + buffer = io.BytesIO() + img.save(buffer, format='PNG', optimize=True) + png_data = buffer.getvalue() + + if len(png_data) <= max_size_bytes: + logger.info(f"Image optimized as PNG: {len(png_data) / (1024*1024):.2f}MB") + return png_data + + # PNG too large, try JPEG with high quality + for quality in [95, 90, 85, 80]: + buffer = io.BytesIO() + img.save(buffer, format='JPEG', quality=quality, optimize=True) + jpeg_data = buffer.getvalue() + + if len(jpeg_data) <= max_size_bytes: + size_mb = len(jpeg_data) / (1024 * 1024) + logger.info(f"Image optimized as JPEG (quality {quality}): {size_mb:.2f}MB") + return jpeg_data + + # Still too large, try resizing (preserve aspect ratio) + logger.warning("Image still too large, attempting resize...") + + scale = 0.95 + while scale > 0.5: + new_width = int(img.width * scale) + new_height = int(img.height * scale) + + resized = img.resize((new_width, new_height), Image.Resampling.LANCZOS) + + buffer = io.BytesIO() + resized.save(buffer, format='JPEG', quality=85, optimize=True) + data = buffer.getvalue() + + if len(data) <= max_size_bytes: + size_mb = len(data) / (1024 * 1024) + logger.warning(f"Image resized to {new_width}x{new_height} ({scale*100:.0f}%): {size_mb:.2f}MB") + return data + + scale -= 0.05 + + raise ValueError(f"Cannot optimize image to under {max_size_mb}MB") + + except Exception as e: + logger.error(f"Error optimizing image: {e}") + raise + + +def convert_image_to_base64(image_path: str) -> str: + """ + Convert image file to base64 string, optimizing if necessary + + Args: + image_path: Path to image file + + Returns: + str: Base64 encoded image data + """ + try: + # Check if optimization is needed + current_size_mb = get_file_size_mb(image_path) + + if current_size_mb <= 20.0: + # Read directly if under limit + with open(image_path, 'rb') as f: + image_data = f.read() + else: + # Optimize if over limit + image_data = optimize_image_for_flux(image_path) + + return encode_image_base64(image_data) + + except Exception as e: + logger.error(f"Error converting image to base64: {e}") + raise + + +def validate_aspect_ratio(width: int, height: int, target_ratio: str) -> bool: + """ + Validate if image dimensions match target aspect ratio + + Args: + width: Image width + height: Image height + target_ratio: Target aspect ratio (e.g., "16:9", "1:1") + + Returns: + bool: True if matches (within tolerance) + """ + try: + ratio_parts = target_ratio.split(':') + if len(ratio_parts) != 2: + return False + + target_w, target_h = int(ratio_parts[0]), int(ratio_parts[1]) + target_ratio_val = target_w / target_h + actual_ratio_val = width / height + + # Allow 5% tolerance + tolerance = 0.05 + return abs(target_ratio_val - actual_ratio_val) / target_ratio_val <= tolerance + + except Exception: + return False + + +def get_optimal_aspect_ratio(width: int, height: int) -> str: + """ + Get the optimal aspect ratio for given dimensions + + Args: + width: Image width + height: Image height + + Returns: + str: Optimal aspect ratio + """ + common_ratios = { + "1:1": 1.0, + "4:3": 4/3, + "3:4": 3/4, + "16:9": 16/9, + "9:16": 9/16, + "21:9": 21/9, + "9:21": 9/21 + } + + actual_ratio = width / height + + # Find closest ratio + closest_ratio = "1:1" + min_diff = float('inf') + + for ratio_name, ratio_val in common_ratios.items(): + diff = abs(ratio_val - actual_ratio) + if diff < min_diff: + min_diff = diff + closest_ratio = ratio_name + + return closest_ratio diff --git a/src/utils/validation.py b/src/utils/validation.py new file mode 100644 index 0000000..397a437 --- /dev/null +++ b/src/utils/validation.py @@ -0,0 +1,387 @@ +"""Validation utilities for FLUX.1 Edit MCP Server""" + +import re +import logging +from typing import Dict, Any, List, Tuple, Optional + +logger = logging.getLogger(__name__) + + +def validate_edit_parameters(arguments: Dict[str, Any]) -> Tuple[bool, Optional[str]]: + """ + Validate parameters for flux_edit_image + + Args: + arguments: Tool arguments + + Returns: + tuple: (is_valid, error_message) + """ + try: + # Check required parameters + if 'input_image_b64' not in arguments: + return False, "input_image_b64 is required" + + if 'prompt' not in arguments: + return False, "prompt is required" + + if 'seed' not in arguments: + return False, "seed is required" + + # Validate input_image_b64 + input_image_b64 = arguments['input_image_b64'] + if not isinstance(input_image_b64, str) or not input_image_b64.strip(): + return False, "input_image_b64 must be a non-empty string" + + # Basic base64 format check + try: + # Remove data URL prefix if present + if input_image_b64.startswith('data:'): + comma_index = input_image_b64.find(',') + if comma_index != -1: + b64_data = input_image_b64[comma_index + 1:] + else: + return False, "Invalid data URL format in input_image_b64" + else: + b64_data = input_image_b64 + + # Check base64 format (basic validation) + if not re.match(r'^[A-Za-z0-9+/]*={0,2}$', b64_data.replace('\n', '').replace('\r', '')): + return False, "Invalid base64 format in input_image_b64" + + except Exception as e: + return False, f"Error validating base64 data: {str(e)}" + + # Validate prompt + prompt = arguments['prompt'] + if not isinstance(prompt, str) or not prompt.strip(): + return False, "prompt must be a non-empty string" + + # Check prompt length (reasonable limit) + if len(prompt) > 10000: # FLUX doesn't specify limit, but this is reasonable + return False, "prompt is too long (max: 10000 characters)" + + # Validate seed + seed = arguments['seed'] + if not isinstance(seed, int): + return False, "seed must be an integer" + + 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)}" + + if 'save_to_file' in arguments: + save_to_file = arguments['save_to_file'] + if not isinstance(save_to_file, bool): + return False, "save_to_file must be a boolean" + + return True, None + + except Exception as e: + logger.error(f"Error validating edit parameters: {e}") + return False, f"Validation error: {str(e)}" + + +def validate_file_parameters(arguments: Dict[str, Any]) -> Tuple[bool, Optional[str]]: + """ + Validate parameters for file-based editing functions + + Args: + arguments: Tool arguments + + Returns: + tuple: (is_valid, error_message) + """ + try: + # Check required parameters + if 'input_image_name' not in arguments: + return False, "input_image_name is required" + + if 'prompt' not in arguments: + return False, "prompt is required" + + if 'seed' not in arguments: + return False, "seed is required" + + # Validate input_image_name + input_image_name = arguments['input_image_name'] + if not isinstance(input_image_name, str) or not input_image_name.strip(): + return False, "input_image_name must be a non-empty string" + + # Check for path traversal attempts + if '..' in input_image_name or '/' in input_image_name or '\\' in input_image_name: + return False, "input_image_name cannot contain path separators or '..' for security" + + # Check file extension (reasonable image formats) + valid_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp', '.tiff', '.tif'] + if not any(input_image_name.lower().endswith(ext) for ext in valid_extensions): + return False, f"input_image_name must have a valid image extension: {', '.join(valid_extensions)}" + + # Validate prompt + prompt = arguments['prompt'] + if not isinstance(prompt, str) or not prompt.strip(): + return False, "prompt must be a non-empty string" + + # Check prompt length + if len(prompt) > 10000: + return False, "prompt is too long (max: 10000 characters)" + + # Validate seed + seed = arguments['seed'] + if not isinstance(seed, int): + return False, "seed must be an integer" + + 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)}" + + if 'save_to_file' in arguments: + save_to_file = arguments['save_to_file'] + if not isinstance(save_to_file, bool): + return False, "save_to_file must be a boolean" + + return True, None + + except Exception as e: + logger.error(f"Error validating file parameters: {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]]: + """ + Validate parameters for validate_image tool + + Args: + arguments: Tool arguments + + Returns: + tuple: (is_valid, error_message) + """ + try: + # Check required parameters + if 'image_path' not in arguments: + return False, "image_path is required" + + # Validate image_path + image_path = arguments['image_path'] + if not isinstance(image_path, str) or not image_path.strip(): + return False, "image_path must be a non-empty string" + + return True, None + + except Exception as e: + logger.error(f"Error validating image path parameter: {e}") + return False, f"Validation error: {str(e)}" + + +def sanitize_prompt(prompt: str) -> str: + """ + Sanitize prompt text (basic cleanup) + + Args: + prompt: Raw prompt text + + Returns: + str: Sanitized prompt + """ + try: + # Remove excessive whitespace + sanitized = ' '.join(prompt.strip().split()) + + # Remove or replace potentially problematic characters + # Note: FLUX.1 Kontext is quite robust, so minimal sanitization needed + + # Remove null bytes + sanitized = sanitized.replace('\x00', '') + + # Limit length if too long + if len(sanitized) > 10000: + sanitized = sanitized[:10000] + logger.warning("Prompt truncated to 10000 characters") + + return sanitized + + except Exception as e: + logger.error(f"Error sanitizing prompt: {e}") + return prompt # Return original if sanitization fails + + +def validate_aspect_ratio_format(aspect_ratio: str) -> bool: + """ + Validate aspect ratio format (e.g., "16:9", "1:1") + + Args: + aspect_ratio: Aspect ratio string + + Returns: + bool: True if valid format + """ + try: + if ':' not in aspect_ratio: + return False + + parts = aspect_ratio.split(':') + if len(parts) != 2: + return False + + # Check if both parts are positive integers + width_ratio = int(parts[0]) + height_ratio = int(parts[1]) + + return width_ratio > 0 and height_ratio > 0 + + except ValueError: + return False + + +def validate_seed_range(seed: int) -> bool: + """ + Validate seed is within valid range + + Args: + seed: Seed value + + Returns: + bool: True if valid + """ + return isinstance(seed, int) and 0 <= seed <= 2**32 - 1 + + +def validate_base64_image(base64_str: str) -> Tuple[bool, Optional[str]]: + """ + Validate base64 image data + + Args: + base64_str: Base64 encoded image + + Returns: + tuple: (is_valid, error_message) + """ + try: + from .image_utils import decode_image_base64, get_image_dimensions_from_bytes + + # Try to decode + image_data = decode_image_base64(base64_str) + + # Check size (20MB limit for FLUX) + size_mb = len(image_data) / (1024 * 1024) + if size_mb > 20: + return False, f"Image size {size_mb:.2f}MB exceeds 20MB limit" + + # Try to get dimensions (validates it's a real image) + width, height = get_image_dimensions_from_bytes(image_data) + if width == 0 or height == 0: + return False, "Invalid image data - cannot determine dimensions" + + # Check reasonable dimension limits + if width > 8192 or height > 8192: + return False, f"Image dimensions {width}x{height} too large (max: 8192x8192)" + + if width < 32 or height < 32: + return False, f"Image dimensions {width}x{height} too small (min: 32x32)" + + return True, None + + except Exception as e: + return False, f"Invalid base64 image data: {str(e)}" + + +def validate_filename_safety(filename: str) -> bool: + """ + Check if filename is safe (no path traversal, reasonable characters) + + Args: + filename: Filename to check + + Returns: + bool: True if safe + """ + try: + # Check for path traversal + if '..' in filename or '/' in filename or '\\' in filename: + return False + + # Check for reserved names (Windows) + reserved = ['CON', 'PRN', 'AUX', 'NUL'] + [f'COM{i}' for i in range(1, 10)] + [f'LPT{i}' for i in range(1, 10)] + name_without_ext = filename.split('.')[0].upper() + if name_without_ext in reserved: + return False + + # Check for invalid characters + invalid_chars = '<>:"|?*' + if any(char in filename for char in invalid_chars): + return False + + # Check length + if len(filename) > 255: + return False + + return True + + except Exception: + return False diff --git a/start.bat b/start.bat new file mode 100644 index 0000000..07c3509 --- /dev/null +++ b/start.bat @@ -0,0 +1,25 @@ +@echo off +echo FLUX.1 Edit MCP Server - Quick Start +echo ==================================== + +echo Checking dependencies first... +REM Set UTF-8 encoding to prevent Unicode errors +chcp 65001 >nul 2>&1 +set PYTHONIOENCODING=utf-8 +set PYTHONUTF8=1 +python check_dependencies.py +if errorlevel 1 ( + echo. + echo Dependencies check failed. Running installation... + call install_dependencies.bat + if errorlevel 1 ( + echo Failed to install dependencies. + pause + exit /b 1 + ) +) + +echo. +echo Dependencies OK! Running original startup script... +echo. +call run.bat diff --git a/test_fixes.py b/test_fixes.py new file mode 100644 index 0000000..f0261b1 --- /dev/null +++ b/test_fixes.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +""" +Test script to verify FLUX.1 Edit MCP Server fixes +This script tests if the Unicode and JSON parsing issues are resolved +""" + +import sys +import os +from pathlib import Path + +# Add src to path for imports +sys.path.insert(0, str(Path(__file__).parent / 'src')) + +def test_console_encoding(): + """Test console encoding with safe characters""" + print("Testing console output with safe characters...") + print("[OK] ASCII characters work fine") + print("[ERROR] Error messages use brackets instead of Unicode") + print("[SUCCESS] Success messages use brackets instead of Unicode") + print("[INFO] Info messages work properly") + return True + +def test_dependency_imports(): + """Test importing dependencies silently""" + print("Testing dependency imports...") + + missing_deps = [] + + try: + import aiohttp + # Silent check + except ImportError: + missing_deps.append("aiohttp") + + try: + import mcp + # Silent check + except ImportError: + missing_deps.append("mcp") + + try: + from src.connector import Config + # Silent check + except ImportError: + missing_deps.append("src.connector.Config") + + try: + from src.server import main + # Silent check + except ImportError: + missing_deps.append("src.server.main") + + if missing_deps: + print(f"[ERROR] Missing dependencies: {', '.join(missing_deps)}") + return False + else: + print("[SUCCESS] All imports successful") + return True + +def test_server_creation(): + """Test creating server instance without starting it""" + print("Testing server creation...") + + try: + from src.server import create_server + server = create_server() + print("[SUCCESS] Server instance created successfully") + return True + except Exception as e: + print(f"[ERROR] Server creation failed: {e}") + return False + +def main(): + """Main test function""" + print("FLUX.1 Edit MCP Server - Fix Verification") + print("=" * 50) + + # Set UTF-8 encoding environment variables + os.environ['PYTHONIOENCODING'] = 'utf-8' + os.environ['PYTHONUTF8'] = '1' + + tests = [ + ("Console encoding", test_console_encoding), + ("Dependency imports", test_dependency_imports), + ("Server creation", test_server_creation), + ] + + passed = 0 + total = len(tests) + + for test_name, test_func in tests: + print(f"\nRunning test: {test_name}") + print("-" * 30) + + try: + if test_func(): + passed += 1 + print(f"[PASSED] {test_name}") + else: + print(f"[FAILED] {test_name}") + except Exception as e: + print(f"[FAILED] {test_name}: {e}") + + print("\n" + "=" * 50) + print(f"Test Results: {passed}/{total} tests passed") + + if passed == total: + print("[SUCCESS] All tests passed! The fixes should work.") + print("\nTo run the server:") + print("1. Make sure your .env file is configured") + print("2. Run: start.bat or run.bat") + print("3. The server should start without Unicode errors") + return 0 + else: + print(f"[FAILED] {total - passed} tests failed. Check the errors above.") + return 1 + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/run_tests.py b/tests/run_tests.py new file mode 100644 index 0000000..b2a2988 --- /dev/null +++ b/tests/run_tests.py @@ -0,0 +1,106 @@ +"""Test runner for all FLUX.1 Edit MCP Server tests""" + +import unittest +import sys +from pathlib import Path + +# Add src to path +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +# Import all test modules +from test_config import TestConfig +from test_image_utils import TestImageUtils +from test_validation import TestValidation +from test_flux_client import TestFluxEditClient, TestFluxEditClientAsync +from test_handlers import TestToolHandlers + + +def create_test_suite(): + """Create and return the complete test suite""" + suite = unittest.TestSuite() + + # Add all test cases + suite.addTest(unittest.makeSuite(TestConfig)) + suite.addTest(unittest.makeSuite(TestImageUtils)) + suite.addTest(unittest.makeSuite(TestValidation)) + suite.addTest(unittest.makeSuite(TestFluxEditClient)) + suite.addTest(unittest.makeSuite(TestFluxEditClientAsync)) + suite.addTest(unittest.makeSuite(TestToolHandlers)) + + return suite + + +def run_tests(): + """Run all tests and return results""" + # Setup test runner + runner = unittest.TextTestRunner( + verbosity=2, + stream=sys.stdout, + descriptions=True, + failfast=False + ) + + # Create and run test suite + suite = create_test_suite() + result = runner.run(suite) + + # Print summary + print(f"\n{'='*70}") + print("TEST SUMMARY") + print(f"{'='*70}") + print(f"Tests run: {result.testsRun}") + print(f"Failures: {len(result.failures)}") + print(f"Errors: {len(result.errors)}") + print(f"Skipped: {len(result.skipped) if hasattr(result, 'skipped') else 0}") + + if result.failures: + print(f"\nFAILURES:") + for test, traceback in result.failures: + print(f" - {test}: {traceback.split('AssertionError: ')[-1].split()[0] if 'AssertionError:' in traceback else 'Unknown failure'}") + + if result.errors: + print(f"\nERRORS:") + for test, traceback in result.errors: + error_msg = traceback.split('\n')[-2] if traceback.split('\n')[-2] else 'Unknown error' + print(f" - {test}: {error_msg}") + + success = len(result.failures) == 0 and len(result.errors) == 0 + print(f"\nOVERALL: {'✅ PASSED' if success else '❌ FAILED'}") + + return success + + +def run_specific_test(test_name: str): + """Run a specific test module""" + test_modules = { + 'config': TestConfig, + 'image_utils': TestImageUtils, + 'validation': TestValidation, + 'flux_client': TestFluxEditClient, + 'flux_client_async': TestFluxEditClientAsync, + 'handlers': TestToolHandlers + } + + if test_name not in test_modules: + print(f"❌ Unknown test module: {test_name}") + print(f"Available modules: {', '.join(test_modules.keys())}") + return False + + suite = unittest.makeSuite(test_modules[test_name]) + runner = unittest.TextTestRunner(verbosity=2) + result = runner.run(suite) + + return len(result.failures) == 0 and len(result.errors) == 0 + + +if __name__ == '__main__': + # Check for specific test argument + if len(sys.argv) > 1: + test_name = sys.argv[1] + success = run_specific_test(test_name) + else: + # Run all tests + success = run_tests() + + # Exit with appropriate code + sys.exit(0 if success else 1) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..29d4df6 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,159 @@ +"""Unit tests for Config class""" + +import unittest +import tempfile +import os +from pathlib import Path +from unittest.mock import patch, MagicMock + +# Add src to path +import sys +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from src.connector.config import Config + + +class TestConfig(unittest.TestCase): + """Test cases for Config class""" + + def setUp(self): + """Set up test fixtures""" + self.temp_dir = tempfile.mkdtemp() + self.temp_path = Path(self.temp_dir) + + def tearDown(self): + """Clean up test fixtures""" + import shutil + if self.temp_path.exists(): + shutil.rmtree(self.temp_path) + + @patch.dict(os.environ, { + 'FLUX_API_KEY': 'test_api_key_12345', + 'LOG_LEVEL': 'DEBUG', + 'MAX_IMAGE_SIZE_MB': '25', + 'DEFAULT_TIMEOUT': '600' + }) + def test_config_initialization(self): + """Test config initialization with environment variables""" + config = Config() + + self.assertEqual(config.api_key, 'test_api_key_12345') + self.assertEqual(config.log_level, 'DEBUG') + self.assertEqual(config.max_image_size_mb, 25) + self.assertEqual(config.default_timeout, 600) + + def test_config_defaults(self): + """Test config defaults when env vars not set""" + # Clear environment + env_vars_to_clear = [ + 'FLUX_API_KEY', 'LOG_LEVEL', 'MAX_IMAGE_SIZE_MB', + 'DEFAULT_TIMEOUT', 'POLLING_INTERVAL', 'MAX_POLLING_ATTEMPTS' + ] + + with patch.dict(os.environ, {}, clear=True): + config = Config() + + self.assertEqual(config.api_key, '') + self.assertEqual(config.log_level, 'INFO') + self.assertEqual(config.max_image_size_mb, 20) + self.assertEqual(config.default_timeout, 300) + self.assertEqual(config.polling_interval, 2) + self.assertEqual(config.max_polling_attempts, 150) + + def test_api_url_generation(self): + """Test API URL generation""" + config = Config() + + edit_url = config.get_api_url(config.EDIT_ENDPOINT) + result_url = config.get_api_url(config.RESULT_ENDPOINT) + + expected_edit = f"{config.api_base_url}/flux-kontext-pro" + expected_result = f"{config.api_base_url}/v1/get_result" + + self.assertEqual(edit_url, expected_edit) + self.assertEqual(result_url, expected_result) + + def test_filename_generation(self): + """Test filename generation""" + config = Config() + + base_name = "fluxedit_123456_20250826_143022" + + # Test different file numbers and extensions + filename_000 = config.generate_filename(base_name, 0, 'png') + filename_001 = config.generate_filename(base_name, 1, 'png') + filename_json = config.generate_filename(base_name, 1, 'json') + + self.assertEqual(filename_000, "fluxedit_123456_20250826_143022_000.png") + self.assertEqual(filename_001, "fluxedit_123456_20250826_143022_001.png") + self.assertEqual(filename_json, "fluxedit_123456_20250826_143022_001.json") + + def test_base_name_generation(self): + """Test base name generation""" + config = Config() + + # Test with seed + base_name_with_seed = config.generate_base_name(12345) + self.assertIn("fluxedit_12345_", base_name_with_seed) + + # Test simple generation + base_name_simple = config.generate_base_name_simple() + self.assertIn("fluxedit_", base_name_simple) + self.assertNotIn("_12345_", base_name_simple) # No seed in simple + + @patch.dict(os.environ, { + 'FLUX_API_KEY': 'valid_key', + 'MAX_IMAGE_SIZE_MB': '20', + 'DEFAULT_TIMEOUT': '300' + }) + def test_validation_success(self): + """Test successful validation""" + config = Config() + self.assertTrue(config.validate()) + + def test_validation_failures(self): + """Test validation failures""" + # Test missing API key + with patch.dict(os.environ, {'FLUX_API_KEY': ''}, clear=True): + config = Config() + self.assertFalse(config.validate()) + + # Test invalid image size + with patch.dict(os.environ, { + 'FLUX_API_KEY': 'valid_key', + 'MAX_IMAGE_SIZE_MB': '0' + }, clear=True): + config = Config() + self.assertFalse(config.validate()) + + # Test invalid timeout + with patch.dict(os.environ, { + 'FLUX_API_KEY': 'valid_key', + 'DEFAULT_TIMEOUT': '-1' + }, clear=True): + config = Config() + self.assertFalse(config.validate()) + + def test_max_image_size_bytes(self): + """Test max image size in bytes calculation""" + config = Config() + config.max_image_size_mb = 20 + + expected_bytes = 20 * 1024 * 1024 + self.assertEqual(config.get_max_image_size_bytes(), expected_bytes) + + @patch('pathlib.Path.mkdir') + @patch('pathlib.Path.exists') + def test_directory_creation(self, mock_exists, mock_mkdir): + """Test directory creation logic""" + mock_exists.return_value = True + + # This should not raise an exception + config = Config() + + # Verify mkdir was called + mock_mkdir.assert_called() + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_flux_client.py b/tests/test_flux_client.py new file mode 100644 index 0000000..18b15af --- /dev/null +++ b/tests/test_flux_client.py @@ -0,0 +1,308 @@ +"""Unit tests for FLUX API client""" + +import unittest +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch +from pathlib import Path + +# Add src to path +import sys +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from src.connector.flux_client import FluxEditClient, FluxEditRequest, FluxEditResponse +from src.connector.config import Config + + +class TestFluxEditClient(unittest.TestCase): + """Test cases for FLUX API client""" + + def setUp(self): + """Set up test fixtures""" + # Mock config + self.config = MagicMock(spec=Config) + self.config.api_key = 'test_api_key' + self.config.default_timeout = 30 + self.config.polling_interval = 1 + self.config.max_polling_attempts = 5 + self.config.get_api_url.side_effect = lambda endpoint: f"https://api.test.com{endpoint}" + + self.client = FluxEditClient(self.config) + + # Sample request + self.sample_request = FluxEditRequest( + input_image_b64='test_base64_data', + prompt='Make the sky blue', + seed=12345, + aspect_ratio='16:9' + ) + + def tearDown(self): + """Clean up test fixtures""" + # Ensure client session is closed + if hasattr(self.client, 'session') and self.client.session: + asyncio.create_task(self.client.close()) + + @patch('aiohttp.ClientSession') + async def test_create_edit_request_success(self, mock_session_class): + """Test successful edit request creation""" + # Mock session and response + mock_session = AsyncMock() + mock_response = AsyncMock() + mock_response.status = 200 + mock_response.json.return_value = {'id': 'request_12345'} + + mock_session.post.return_value.__aenter__.return_value = mock_response + mock_session_class.return_value = mock_session + + # Test + request_id = await self.client._create_edit_request(self.sample_request) + + # Verify + self.assertEqual(request_id, 'request_12345') + mock_session.post.assert_called_once() + + # Check payload structure + call_args = mock_session.post.call_args + payload = call_args[1]['json'] + self.assertEqual(payload['prompt'], 'Make the sky blue') + self.assertEqual(payload['seed'], 12345) + self.assertEqual(payload['input_image'], 'test_base64_data') + + @patch('aiohttp.ClientSession') + async def test_create_edit_request_failure(self, mock_session_class): + """Test edit request creation failure""" + # Mock session and response + mock_session = AsyncMock() + mock_response = AsyncMock() + mock_response.status = 400 + mock_response.text.return_value = 'Bad Request' + + mock_session.post.return_value.__aenter__.return_value = mock_response + mock_session_class.return_value = mock_session + + # Test + request_id = await self.client._create_edit_request(self.sample_request) + + # Verify + self.assertIsNone(request_id) + + @patch('aiohttp.ClientSession') + async def test_poll_result_success(self, mock_session_class): + """Test successful result polling""" + # Mock session and responses + mock_session = AsyncMock() + + # First response: processing + mock_response_processing = AsyncMock() + mock_response_processing.status = 200 + mock_response_processing.json.return_value = {'status': 'processing'} + + # Second response: ready + mock_response_ready = AsyncMock() + mock_response_ready.status = 200 + mock_response_ready.json.return_value = { + 'status': 'ready', + 'result': {'sample': 'https://example.com/image.png'} + } + + # Mock to return processing first, then ready + mock_session.get.return_value.__aenter__.side_effect = [ + mock_response_processing, + mock_response_ready + ] + mock_session_class.return_value = mock_session + + # Test + result = await self.client._poll_result('test_request_id') + + # Verify + self.assertIsNotNone(result) + self.assertEqual(result['status'], 'ready') + self.assertIn('result', result) + self.assertEqual(mock_session.get.call_count, 2) + + @patch('aiohttp.ClientSession') + async def test_poll_result_timeout(self, mock_session_class): + """Test polling timeout""" + # Mock session to always return processing + mock_session = AsyncMock() + mock_response = AsyncMock() + mock_response.status = 200 + mock_response.json.return_value = {'status': 'processing'} + + mock_session.get.return_value.__aenter__.return_value = mock_response + mock_session_class.return_value = mock_session + + # Test + result = await self.client._poll_result('test_request_id') + + # Verify - should timeout after max attempts + self.assertIsNone(result) + self.assertEqual(mock_session.get.call_count, self.config.max_polling_attempts) + + @patch('aiohttp.ClientSession') + async def test_download_result_image_success(self, mock_session_class): + """Test successful image download""" + # Mock session and response + mock_session = AsyncMock() + mock_response = AsyncMock() + mock_response.status = 200 + mock_response.read.return_value = b'fake_image_data' + + mock_session.get.return_value.__aenter__.return_value = mock_response + mock_session_class.return_value = mock_session + + # Test + image_data = await self.client._download_result_image('https://example.com/image.png') + + # Verify + self.assertEqual(image_data, b'fake_image_data') + mock_session.get.assert_called_once_with('https://example.com/image.png') + + @patch('aiohttp.ClientSession') + async def test_download_result_image_failure(self, mock_session_class): + """Test image download failure""" + # Mock session and response + mock_session = AsyncMock() + mock_response = AsyncMock() + mock_response.status = 404 + + mock_session.get.return_value.__aenter__.return_value = mock_response + mock_session_class.return_value = mock_session + + # Test + image_data = await self.client._download_result_image('https://example.com/image.png') + + # Verify + self.assertIsNone(image_data) + + def test_get_image_size(self): + """Test image size detection from bytes""" + # Create a small test image in memory + from PIL import Image + import io + + # Create 10x20 test image + img = Image.new('RGB', (10, 20), color='red') + buffer = io.BytesIO() + img.save(buffer, format='PNG') + image_data = buffer.getvalue() + + # Test + size = self.client._get_image_size(image_data) + + # Verify + self.assertEqual(size, (10, 20)) + + def test_get_image_size_invalid_data(self): + """Test image size detection with invalid data""" + size = self.client._get_image_size(b'invalid_image_data') + self.assertIsNone(size) + + @patch.object(FluxEditClient, '_create_edit_request') + @patch.object(FluxEditClient, '_poll_result') + @patch.object(FluxEditClient, '_download_result_image') + async def test_edit_image_success(self, mock_download, mock_poll, mock_create): + """Test complete successful edit flow""" + # Setup mocks + mock_create.return_value = 'request_123' + mock_poll.return_value = { + 'status': 'ready', + 'result': {'sample': 'https://example.com/result.png'} + } + mock_download.return_value = b'edited_image_data' + + # Test + response = await self.client.edit_image(self.sample_request) + + # Verify + self.assertTrue(response.success) + self.assertEqual(response.edited_image_data, b'edited_image_data') + self.assertEqual(response.request_id, 'request_123') + self.assertEqual(response.result_url, 'https://example.com/result.png') + self.assertGreater(response.execution_time, 0) + + @patch.object(FluxEditClient, '_create_edit_request') + async def test_edit_image_create_failure(self, mock_create): + """Test edit flow with creation failure""" + # Setup mock + mock_create.return_value = None + + # Test + response = await self.client.edit_image(self.sample_request) + + # Verify + self.assertFalse(response.success) + self.assertIn('Failed to create edit request', response.error_message) + + @patch.object(FluxEditClient, '_create_edit_request') + @patch.object(FluxEditClient, '_poll_result') + async def test_edit_image_poll_failure(self, mock_poll, mock_create): + """Test edit flow with polling failure""" + # Setup mocks + mock_create.return_value = 'request_123' + mock_poll.return_value = None + + # Test + response = await self.client.edit_image(self.sample_request) + + # Verify + self.assertFalse(response.success) + self.assertIn('Failed to get edit result', response.error_message) + + async def test_context_manager(self): + """Test async context manager functionality""" + async with FluxEditClient(self.config) as client: + self.assertIsInstance(client, FluxEditClient) + + # Session should be closed after context + if hasattr(client, 'session'): + self.assertTrue(client.session is None or client.session.closed) + + +# Test helper to run async tests +class AsyncTestCase(unittest.TestCase): + def setUp(self): + self.loop = asyncio.new_event_loop() + asyncio.set_event_loop(self.loop) + + def tearDown(self): + self.loop.close() + + def async_test(self, coro): + return self.loop.run_until_complete(coro) + + +class TestFluxEditClientAsync(AsyncTestCase): + """Async test runner for FLUX client tests""" + + def setUp(self): + super().setUp() + self.config = MagicMock(spec=Config) + self.config.api_key = 'test_api_key' + self.config.default_timeout = 30 + self.config.polling_interval = 0.1 # Faster for tests + self.config.max_polling_attempts = 3 + self.config.get_api_url.side_effect = lambda endpoint: f"https://api.test.com{endpoint}" + + self.client = FluxEditClient(self.config) + self.sample_request = FluxEditRequest( + input_image_b64='test_base64_data', + prompt='Make the sky blue', + seed=12345, + aspect_ratio='16:9' + ) + + def test_async_context_manager(self): + """Test async context manager""" + async def run_test(): + async with FluxEditClient(self.config) as client: + self.assertIsInstance(client, FluxEditClient) + return True + + result = self.async_test(run_test()) + self.assertTrue(result) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_handlers.py b/tests/test_handlers.py new file mode 100644 index 0000000..eecf77c --- /dev/null +++ b/tests/test_handlers.py @@ -0,0 +1,360 @@ +"""Unit tests for MCP tool handlers""" + +import unittest +import tempfile +import asyncio +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch +from PIL import Image +import io + +# Add src to path +import sys +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from src.server.handlers import ToolHandlers +from src.connector.config import Config +from src.connector.flux_client import FluxEditResponse +from mcp.types import TextContent, ImageContent + + +class AsyncTestCase(unittest.TestCase): + """Base class for async tests""" + + def setUp(self): + self.loop = asyncio.new_event_loop() + asyncio.set_event_loop(self.loop) + + def tearDown(self): + self.loop.close() + + def async_test(self, coro): + return self.loop.run_until_complete(coro) + + +class TestToolHandlers(AsyncTestCase): + """Test cases for MCP tool handlers""" + + def setUp(self): + super().setUp() + + # Create temporary directory + self.temp_dir = tempfile.mkdtemp() + self.temp_path = Path(self.temp_dir) + + # Mock config + self.config = MagicMock(spec=Config) + self.config.base_path = self.temp_path + self.config.input_path = self.temp_path / 'input_images' + self.config.generated_images_path = self.temp_path / 'generated_images' + self.config.max_image_size_mb = 20 + self.config.default_aspect_ratio = '1:1' + self.config.safety_tolerance = 2 + self.config.OUTPUT_FORMAT = 'png' + self.config.prompt_upsampling = False + self.config.MODEL_NAME = 'flux-kontext-pro' + self.config.save_parameters = True + + # Setup directories + self.config.input_path.mkdir(parents=True, exist_ok=True) + self.config.generated_images_path.mkdir(parents=True, exist_ok=True) + + # Mock config methods + self.config.ensure_output_directory.return_value = None + self.config.generate_base_name.return_value = 'fluxedit_12345_20250826_143022' + self.config.generate_filename.side_effect = lambda base, num, ext: f'{base}_{num:03d}.{ext}' + self.config.get_output_path.side_effect = lambda base, num, ext: self.config.generated_images_path / f'{base}_{num:03d}.{ext}' + + # Create test image + self.test_image = Image.new('RGB', (100, 100), color='blue') + buffer = io.BytesIO() + self.test_image.save(buffer, format='PNG') + self.test_image_data = buffer.getvalue() + self.test_image_b64 = self._encode_image_b64(self.test_image_data) + + # Create test image file + self.test_image_file = self.config.input_path / 'test.png' + self.test_image.save(self.test_image_file) + + # Initialize handlers + self.handlers = ToolHandlers(self.config) + + def tearDown(self): + """Clean up test fixtures""" + import shutil + super().tearDown() + if self.temp_path.exists(): + shutil.rmtree(self.temp_path) + + def _encode_image_b64(self, image_data: bytes) -> str: + """Helper to encode image as base64""" + import base64 + return base64.b64encode(image_data).decode('utf-8') + + def _create_mock_flux_response(self, success: bool = True) -> FluxEditResponse: + """Helper to create mock FLUX response""" + if success: + return FluxEditResponse( + success=True, + edited_image_data=self.test_image_data, + image_size=(100, 100), + execution_time=5.5, + request_id='test_request_123', + result_url='https://example.com/result.png', + metadata={'seed': 12345} + ) + else: + return FluxEditResponse( + success=False, + error_message='FLUX edit failed', + execution_time=2.0 + ) + + def test_flux_edit_image_parameter_validation_failure(self): + """Test flux_edit_image with invalid parameters""" + async def run_test(): + # Missing required parameter + arguments = { + 'prompt': 'Make it blue', + 'seed': 12345 + # Missing input_image_b64 + } + + result = await self.handlers.handle_flux_edit_image(arguments) + + self.assertEqual(len(result), 1) + self.assertIsInstance(result[0], TextContent) + self.assertIn('Parameter validation failed', result[0].text) + self.assertIn('input_image_b64 is required', result[0].text) + + self.async_test(run_test()) + + @patch('src.server.handlers.FluxEditClient') + def test_flux_edit_image_success(self, mock_client_class): + """Test successful flux_edit_image""" + async def run_test(): + # Setup mock client + mock_client = AsyncMock() + mock_client_class.return_value.__aenter__.return_value = mock_client + mock_client.edit_image.return_value = self._create_mock_flux_response(success=True) + + # Valid arguments + arguments = { + 'input_image_b64': self.test_image_b64, + 'prompt': 'Make the sky blue', + 'seed': 12345, + 'aspect_ratio': '16:9', + 'save_to_file': True + } + + result = await self.handlers.handle_flux_edit_image(arguments) + + # Verify result structure + self.assertGreater(len(result), 0) + self.assertIsInstance(result[0], TextContent) + self.assertIn('✅ Image edited successfully', result[0].text) + self.assertIn('Seed: 12345', result[0].text) + + # Should have image preview + if len(result) > 1: + self.assertIsInstance(result[1], ImageContent) + + self.async_test(run_test()) + + @patch('src.server.handlers.FluxEditClient') + def test_flux_edit_image_flux_failure(self, mock_client_class): + """Test flux_edit_image with FLUX API failure""" + async def run_test(): + # Setup mock client to return failure + mock_client = AsyncMock() + mock_client_class.return_value.__aenter__.return_value = mock_client + mock_client.edit_image.return_value = self._create_mock_flux_response(success=False) + + arguments = { + 'input_image_b64': self.test_image_b64, + 'prompt': 'Make it blue', + 'seed': 12345 + } + + result = await self.handlers.handle_flux_edit_image(arguments) + + self.assertEqual(len(result), 1) + self.assertIsInstance(result[0], TextContent) + self.assertIn('❌ FLUX edit failed', result[0].text) + + self.async_test(run_test()) + + def test_flux_edit_image_from_file_not_found(self): + """Test flux_edit_image_from_file with non-existent file""" + async def run_test(): + arguments = { + 'input_image_name': 'nonexistent.png', + 'prompt': 'Edit this', + 'seed': 12345 + } + + result = await self.handlers.handle_flux_edit_image_from_file(arguments) + + self.assertEqual(len(result), 1) + self.assertIsInstance(result[0], TextContent) + self.assertIn('❌ File not found', result[0].text) + self.assertIn('nonexistent.png', result[0].text) + + self.async_test(run_test()) + + @patch('src.server.handlers.FluxEditClient') + def test_flux_edit_image_from_file_success(self, mock_client_class): + """Test successful flux_edit_image_from_file""" + async def run_test(): + # Setup mock client + mock_client = AsyncMock() + mock_client_class.return_value.__aenter__.return_value = mock_client + mock_client.edit_image.return_value = self._create_mock_flux_response(success=True) + + arguments = { + 'input_image_name': 'test.png', + 'prompt': 'Make it awesome', + 'seed': 54321, + 'save_to_file': True + } + + result = await self.handlers.handle_flux_edit_image_from_file(arguments) + + # Verify result + self.assertGreater(len(result), 0) + self.assertIsInstance(result[0], TextContent) + self.assertIn('✅ Image edited successfully from file', result[0].text) + self.assertIn('test.png', result[0].text) + self.assertIn('Seed: 54321', result[0].text) + + self.async_test(run_test()) + + def test_validate_image_success(self): + """Test successful image validation""" + async def run_test(): + arguments = {'image_path': str(self.test_image_file)} + + result = await self.handlers.handle_validate_image(arguments) + + self.assertEqual(len(result), 1) + self.assertIsInstance(result[0], TextContent) + self.assertIn('✅ Image validation passed', result[0].text) + self.assertIn('100x100', result[0].text) + + self.async_test(run_test()) + + def test_validate_image_not_found(self): + """Test image validation with non-existent file""" + async def run_test(): + arguments = {'image_path': '/nonexistent/path.png'} + + result = await self.handlers.handle_validate_image(arguments) + + self.assertEqual(len(result), 1) + self.assertIsInstance(result[0], TextContent) + self.assertIn('❌ Image validation failed', result[0].text) + self.assertIn('File not found', result[0].text) + + self.async_test(run_test()) + + def test_move_temp_to_output_success(self): + """Test successful file move operation""" + async def run_test(): + # Create temp directory and file + temp_dir = self.config.base_path / 'temp' + temp_dir.mkdir(exist_ok=True) + temp_file = temp_dir / 'temp_test.png' + self.test_image.save(temp_file) + + arguments = { + 'temp_file_name': 'temp_test.png', + 'output_file_name': 'moved_test.png', + 'copy_only': False + } + + result = await self.handlers.handle_move_temp_to_output(arguments) + + self.assertEqual(len(result), 1) + self.assertIsInstance(result[0], TextContent) + self.assertIn('✅ File moved successfully', result[0].text) + self.assertIn('temp_test.png', result[0].text) + self.assertIn('moved_test.png', result[0].text) + + self.async_test(run_test()) + + def test_move_temp_to_output_not_found(self): + """Test file move with non-existent temp file""" + async def run_test(): + arguments = { + 'temp_file_name': 'nonexistent.png', + 'copy_only': False + } + + result = await self.handlers.handle_move_temp_to_output(arguments) + + self.assertEqual(len(result), 1) + self.assertIsInstance(result[0], TextContent) + self.assertIn('❌ Temp file not found', result[0].text) + + self.async_test(run_test()) + + def test_move_temp_to_output_copy_only(self): + """Test copy-only file operation""" + async def run_test(): + # Create temp directory and file + temp_dir = self.config.base_path / 'temp' + temp_dir.mkdir(exist_ok=True) + temp_file = temp_dir / 'temp_copy.png' + self.test_image.save(temp_file) + + arguments = { + 'temp_file_name': 'temp_copy.png', + 'copy_only': True + } + + result = await self.handlers.handle_move_temp_to_output(arguments) + + self.assertEqual(len(result), 1) + self.assertIsInstance(result[0], TextContent) + self.assertIn('✅ File copied successfully', result[0].text) + + # Original file should still exist + self.assertTrue(temp_file.exists()) + + self.async_test(run_test()) + + def test_seed_management(self): + """Test seed creation and reset functionality""" + # Test seed generation + seed1 = self.handlers._get_or_create_seed() + seed2 = self.handlers._get_or_create_seed() + + # Should return same seed for session + self.assertEqual(seed1, seed2) + + # Test seed reset + self.handlers._reset_seed() + seed3 = self.handlers._get_or_create_seed() + + # Should be different after reset + self.assertNotEqual(seed1, seed3) + + def test_temp_file_operations(self): + """Test temporary file save and move operations""" + # Test saving b64 to temp + filename = 'test_temp.png' + temp_path = self.handlers._save_b64_to_temp_file(self.test_image_b64, filename) + + self.assertTrue(Path(temp_path).exists()) + self.assertIn(filename, temp_path) + + # Test moving to generated images + base_name = 'test_base' + moved_path = self.handlers._move_temp_to_generated(temp_path, base_name, 1) + + self.assertTrue(Path(moved_path).exists()) + self.assertIn('test_base_001.png', moved_path) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_image_utils.py b/tests/test_image_utils.py new file mode 100644 index 0000000..da33b7e --- /dev/null +++ b/tests/test_image_utils.py @@ -0,0 +1,197 @@ +"""Unit tests for image utilities""" + +import unittest +import tempfile +import base64 +from pathlib import Path +from PIL import Image +import io + +# Add src to path +import sys +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from src.utils.image_utils import ( + get_file_size_mb, + validate_image_file, + get_image_dimensions, + get_image_dimensions_from_bytes, + encode_image_base64, + decode_image_base64, + save_image, + get_optimal_aspect_ratio, + validate_aspect_ratio, + convert_image_to_base64 +) + + +class TestImageUtils(unittest.TestCase): + """Test cases for image utility functions""" + + def setUp(self): + """Set up test fixtures""" + self.temp_dir = tempfile.mkdtemp() + self.temp_path = Path(self.temp_dir) + + # Create a test image + self.test_image = Image.new('RGB', (100, 100), color='red') + self.test_image_path = self.temp_path / 'test_image.png' + self.test_image.save(self.test_image_path) + + # Create test image data + buffer = io.BytesIO() + self.test_image.save(buffer, format='PNG') + self.test_image_data = buffer.getvalue() + + def tearDown(self): + """Clean up test fixtures""" + import shutil + if self.temp_path.exists(): + shutil.rmtree(self.temp_path) + + def test_get_file_size_mb(self): + """Test file size calculation""" + size_mb = get_file_size_mb(self.test_image_path) + self.assertGreater(size_mb, 0) + self.assertLess(size_mb, 1) # Small test image should be < 1MB + + def test_validate_image_file_success(self): + """Test successful image validation""" + is_valid, size_mb, error = validate_image_file(str(self.test_image_path), 20) + + self.assertTrue(is_valid) + self.assertGreater(size_mb, 0) + self.assertIsNone(error) + + def test_validate_image_file_not_found(self): + """Test validation of non-existent file""" + is_valid, size_mb, error = validate_image_file('nonexistent.png', 20) + + self.assertFalse(is_valid) + self.assertEqual(size_mb, 0) + self.assertIn('File not found', error) + + def test_validate_image_file_too_large(self): + """Test validation of file too large""" + # Test with very small limit + is_valid, size_mb, error = validate_image_file(str(self.test_image_path), 0.001) + + self.assertFalse(is_valid) + self.assertIn('exceeds', error) + + def test_get_image_dimensions(self): + """Test getting image dimensions""" + width, height = get_image_dimensions(str(self.test_image_path)) + self.assertEqual(width, 100) + self.assertEqual(height, 100) + + def test_get_image_dimensions_from_bytes(self): + """Test getting dimensions from image bytes""" + width, height = get_image_dimensions_from_bytes(self.test_image_data) + self.assertEqual(width, 100) + self.assertEqual(height, 100) + + def test_encode_decode_base64(self): + """Test base64 encoding and decoding""" + # Encode + b64_string = encode_image_base64(self.test_image_data) + self.assertIsInstance(b64_string, str) + + # Decode + decoded_data = decode_image_base64(b64_string) + self.assertEqual(decoded_data, self.test_image_data) + + def test_decode_base64_with_data_url(self): + """Test decoding base64 with data URL prefix""" + b64_string = encode_image_base64(self.test_image_data) + data_url = f"data:image/png;base64,{b64_string}" + + decoded_data = decode_image_base64(data_url) + self.assertEqual(decoded_data, self.test_image_data) + + def test_save_image(self): + """Test saving image data to file""" + output_path = self.temp_path / 'output.png' + + success = save_image(self.test_image_data, str(output_path)) + self.assertTrue(success) + self.assertTrue(output_path.exists()) + + # Verify file content + with open(output_path, 'rb') as f: + saved_data = f.read() + self.assertEqual(saved_data, self.test_image_data) + + def test_get_optimal_aspect_ratio(self): + """Test optimal aspect ratio calculation""" + # Test square image + ratio = get_optimal_aspect_ratio(100, 100) + self.assertEqual(ratio, "1:1") + + # Test wide image + ratio = get_optimal_aspect_ratio(160, 90) + self.assertEqual(ratio, "16:9") + + # Test tall image + ratio = get_optimal_aspect_ratio(90, 160) + self.assertEqual(ratio, "9:16") + + def test_validate_aspect_ratio(self): + """Test aspect ratio validation""" + # Test matching ratio + self.assertTrue(validate_aspect_ratio(100, 100, "1:1")) + self.assertTrue(validate_aspect_ratio(160, 90, "16:9")) + + # Test non-matching ratio (within tolerance) + self.assertTrue(validate_aspect_ratio(161, 90, "16:9")) # Small difference + + # Test non-matching ratio (outside tolerance) + self.assertFalse(validate_aspect_ratio(200, 100, "1:1")) + + def test_convert_image_to_base64(self): + """Test converting image file to base64""" + b64_string = convert_image_to_base64(str(self.test_image_path)) + + self.assertIsInstance(b64_string, str) + + # Verify we can decode it back + decoded_data = decode_image_base64(b64_string) + + # Images should have same dimensions + width, height = get_image_dimensions_from_bytes(decoded_data) + self.assertEqual(width, 100) + self.assertEqual(height, 100) + + def create_large_image_file(self, size_mb: float) -> Path: + """Helper to create a large image file for testing""" + # Calculate dimensions for target size (rough estimate) + # PNG compression varies, so this is approximate + pixels = int((size_mb * 1024 * 1024) / 4) # 4 bytes per pixel (RGBA) + dimension = int(pixels ** 0.5) + + large_image = Image.new('RGBA', (dimension, dimension), color='red') + large_image_path = self.temp_path / 'large_image.png' + large_image.save(large_image_path) + + return large_image_path + + def test_large_image_handling(self): + """Test handling of large images""" + # This test might be slow, so we'll use a smaller "large" image + try: + large_path = self.create_large_image_file(0.1) # 0.1 MB + + # Test validation + is_valid, size_mb, error = validate_image_file(str(large_path), 20) + self.assertTrue(is_valid) + + # Test conversion to base64 + b64_string = convert_image_to_base64(str(large_path)) + self.assertIsInstance(b64_string, str) + + except Exception as e: + self.skipTest(f"Large image test skipped due to: {e}") + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_validation.py b/tests/test_validation.py new file mode 100644 index 0000000..7e30b85 --- /dev/null +++ b/tests/test_validation.py @@ -0,0 +1,218 @@ +"""Unit tests for validation utilities""" + +import unittest +from pathlib import Path + +# Add src to path +import sys +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from src.utils.validation import ( + validate_edit_parameters, + validate_file_parameters, + validate_move_file_parameters, + validate_image_path_parameter, + sanitize_prompt, + validate_aspect_ratio_format, + validate_seed_range, + validate_filename_safety +) + + +class TestValidation(unittest.TestCase): + """Test cases for validation utilities""" + + def test_validate_edit_parameters_success(self): + """Test successful validation of edit parameters""" + valid_args = { + 'input_image_b64': 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==', + 'prompt': 'Make the sky blue', + 'seed': 12345, + 'aspect_ratio': '16:9', + 'save_to_file': True + } + + is_valid, error = validate_edit_parameters(valid_args) + self.assertTrue(is_valid) + self.assertIsNone(error) + + def test_validate_edit_parameters_missing_required(self): + """Test validation fails with missing required parameters""" + # Missing input_image_b64 + args = {'prompt': 'test', 'seed': 12345} + is_valid, error = validate_edit_parameters(args) + self.assertFalse(is_valid) + self.assertIn('input_image_b64 is required', error) + + # Missing prompt + args = {'input_image_b64': 'test_b64', 'seed': 12345} + is_valid, error = validate_edit_parameters(args) + self.assertFalse(is_valid) + self.assertIn('prompt is required', error) + + # Missing seed + args = {'input_image_b64': 'test_b64', 'prompt': 'test'} + is_valid, error = validate_edit_parameters(args) + self.assertFalse(is_valid) + self.assertIn('seed is required', error) + + def test_validate_edit_parameters_invalid_types(self): + """Test validation fails with invalid parameter types""" + # Invalid seed type + args = { + 'input_image_b64': 'valid_b64', + 'prompt': 'test prompt', + 'seed': 'not_a_number' + } + is_valid, error = validate_edit_parameters(args) + self.assertFalse(is_valid) + self.assertIn('seed must be an integer', error) + + # Invalid aspect ratio + args = { + 'input_image_b64': 'valid_b64', + 'prompt': 'test prompt', + 'seed': 12345, + 'aspect_ratio': 'invalid_ratio' + } + is_valid, error = validate_edit_parameters(args) + self.assertFalse(is_valid) + self.assertIn('aspect_ratio must be one of', error) + + def test_validate_edit_parameters_invalid_ranges(self): + """Test validation fails with invalid parameter ranges""" + # Seed out of range + args = { + 'input_image_b64': 'valid_b64', + 'prompt': 'test prompt', + 'seed': -1 + } + is_valid, error = validate_edit_parameters(args) + self.assertFalse(is_valid) + self.assertIn('seed must be between', error) + + # Prompt too long + args = { + 'input_image_b64': 'valid_b64', + 'prompt': 'x' * 10001, # Too long + 'seed': 12345 + } + is_valid, error = validate_edit_parameters(args) + self.assertFalse(is_valid) + self.assertIn('prompt is too long', error) + + def test_validate_file_parameters_success(self): + """Test successful validation of file parameters""" + valid_args = { + 'input_image_name': 'test.png', + 'prompt': 'Edit this image', + 'seed': 54321, + 'aspect_ratio': '1:1', + 'save_to_file': False + } + + is_valid, error = validate_file_parameters(valid_args) + self.assertTrue(is_valid) + self.assertIsNone(error) + + def test_validate_file_parameters_invalid_filename(self): + """Test validation fails with invalid filename""" + # Path traversal attempt + args = { + 'input_image_name': '../../../etc/passwd', + 'prompt': 'test', + 'seed': 12345 + } + is_valid, error = validate_file_parameters(args) + self.assertFalse(is_valid) + self.assertIn('cannot contain path separators', error) + + # Invalid extension + args = { + 'input_image_name': 'test.exe', + 'prompt': 'test', + 'seed': 12345 + } + is_valid, error = validate_file_parameters(args) + self.assertFalse(is_valid) + self.assertIn('must have a valid image extension', error) + + def test_validate_move_file_parameters_success(self): + """Test successful validation of move file parameters""" + valid_args = { + 'temp_file_name': 'temp_image.png', + 'output_file_name': 'output.png', + 'copy_only': True + } + + is_valid, error = validate_move_file_parameters(valid_args) + self.assertTrue(is_valid) + self.assertIsNone(error) + + def test_validate_image_path_parameter_success(self): + """Test successful validation of image path parameter""" + valid_args = {'image_path': '/path/to/image.png'} + + is_valid, error = validate_image_path_parameter(valid_args) + self.assertTrue(is_valid) + self.assertIsNone(error) + + def test_sanitize_prompt(self): + """Test prompt sanitization""" + # Test whitespace normalization + prompt = " This has extra whitespace " + sanitized = sanitize_prompt(prompt) + self.assertEqual(sanitized, "This has extra whitespace") + + # Test null byte removal + prompt = "Test\x00with\x00nulls" + sanitized = sanitize_prompt(prompt) + self.assertEqual(sanitized, "Testwithulls") + + # Test length limiting + long_prompt = "x" * 10001 + sanitized = sanitize_prompt(long_prompt) + self.assertEqual(len(sanitized), 10000) + + def test_validate_aspect_ratio_format(self): + """Test aspect ratio format validation""" + # Valid formats + self.assertTrue(validate_aspect_ratio_format("16:9")) + self.assertTrue(validate_aspect_ratio_format("1:1")) + self.assertTrue(validate_aspect_ratio_format("4:3")) + + # Invalid formats + self.assertFalse(validate_aspect_ratio_format("16-9")) + self.assertFalse(validate_aspect_ratio_format("16:9:1")) + self.assertFalse(validate_aspect_ratio_format("a:b")) + self.assertFalse(validate_aspect_ratio_format("0:1")) + + def test_validate_seed_range(self): + """Test seed range validation""" + # Valid seeds + self.assertTrue(validate_seed_range(0)) + self.assertTrue(validate_seed_range(12345)) + self.assertTrue(validate_seed_range(2**32 - 1)) + + # Invalid seeds + self.assertFalse(validate_seed_range(-1)) + self.assertFalse(validate_seed_range(2**32)) + self.assertFalse(validate_seed_range("not_a_number")) + + def test_validate_filename_safety(self): + """Test filename safety validation""" + # Safe filenames + self.assertTrue(validate_filename_safety("image.png")) + self.assertTrue(validate_filename_safety("my_image_123.jpg")) + self.assertTrue(validate_filename_safety("test-file.png")) + + # Unsafe filenames + self.assertFalse(validate_filename_safety("../image.png")) + self.assertFalse(validate_filename_safety("path/to/image.png")) + self.assertFalse(validate_filename_safety("image<>.png")) + self.assertFalse(validate_filename_safety("CON.png")) # Windows reserved + self.assertFalse(validate_filename_safety("x" * 256)) # Too long + + +if __name__ == '__main__': + unittest.main() diff --git a/troubleshoot.bat b/troubleshoot.bat new file mode 100644 index 0000000..df1def2 --- /dev/null +++ b/troubleshoot.bat @@ -0,0 +1,150 @@ +@echo off +echo FLUX.1 Edit MCP Server - Troubleshooting +echo ======================================= + +echo 1. Checking Python installation... +python --version +if errorlevel 1 ( + echo [ERROR] Python is not installed or not in PATH + goto :end +) else ( + echo [OK] Python is available +) + +echo. +echo 2. Checking pip installation... +pip --version +if errorlevel 1 ( + echo [ERROR] pip is not available + goto :end +) else ( + echo [OK] pip is available +) + +echo. +echo 3. Checking virtual environment... +if exist "venv" ( + echo [OK] Virtual environment directory exists + call venv\Scripts\activate.bat + if errorlevel 1 ( + echo [ERROR] Cannot activate virtual environment + goto :end + ) else ( + echo [OK] Virtual environment activated + ) +) else ( + echo [WARNING] Virtual environment not found + echo Creating virtual environment... + python -m venv venv + if errorlevel 1 ( + echo [ERROR] Failed to create virtual environment + goto :end + ) + call venv\Scripts\activate.bat + echo [OK] Virtual environment created and activated +) + +echo. +echo 4. Checking Python in virtual environment... +where python +python --version + +echo. +echo 5. Checking critical dependencies... +python -c "import aiohttp; print(f'aiohttp: {aiohttp.__version__}')" 2>nul +if errorlevel 1 ( + echo [ERROR] aiohttp not found - installing... + pip install aiohttp==3.11.7 +) else ( + echo [OK] aiohttp is available +) + +python -c "import httpx; print(f'httpx: {httpx.__version__}')" 2>nul +if errorlevel 1 ( + echo [ERROR] httpx not found - installing... + pip install httpx==0.28.1 +) else ( + echo [OK] httpx is available +) + +python -c "import mcp" 2>nul +if errorlevel 1 ( + echo [ERROR] mcp not found - installing... + pip install mcp==1.1.0 +) else ( + echo [OK] mcp is available +) + +python -c "from PIL import Image; print('Pillow: available')" 2>nul +if errorlevel 1 ( + echo [ERROR] Pillow not found - installing... + pip install Pillow==11.0.0 +) else ( + echo [OK] Pillow is available +) + +echo. +echo 6. Checking configuration files... +if exist ".env" ( + echo [OK] .env file exists +) else ( + echo [WARNING] .env file not found + if exist ".env.example" ( + echo Creating .env from example... + copy .env.example .env + echo [OK] .env file created from example + ) else ( + echo [ERROR] .env.example file not found + ) +) + +echo. +echo 7. Checking required directories... +if not exist "input_images" mkdir input_images & echo [OK] Created input_images directory +if not exist "generated_images" mkdir generated_images & echo [OK] Created generated_images directory +if not exist "temp" mkdir temp & echo [OK] Created temp directory + +echo. +echo 8. Testing basic imports... +python -c " +import sys +print(f'Python executable: {sys.executable}') +print(f'Python version: {sys.version}') +print('Testing imports...') +try: + import aiohttp + print('✓ aiohttp imported successfully') +except ImportError as e: + print(f'✗ aiohttp import failed: {e}') + +try: + import mcp + print('✓ mcp imported successfully') +except ImportError as e: + print(f'✗ mcp import failed: {e}') + +try: + from src.connector import Config + print('✓ Local Config imported successfully') +except ImportError as e: + print(f'✗ Local Config import failed: {e}') + +try: + from src.server import main + print('✓ Local server main imported successfully') +except ImportError as e: + print(f'✗ Local server main import failed: {e}') +" + +echo. +echo Troubleshooting complete! +echo. +echo If you still have issues: +echo 1. Delete venv folder and run install_dependencies.bat +echo 2. Make sure you have a stable internet connection +echo 3. Check if your antivirus is blocking Python/pip +echo 4. Try running as administrator +echo. + +:end +pause