initial implementation

This commit is contained in:
2025-08-26 02:35:44 +09:00
commit f1f5ad04f1
33 changed files with 6002 additions and 0 deletions

26
.env.example Normal file
View File

@@ -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

100
.gitignore vendored Normal file
View File

@@ -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

261
CLAUDE.md Normal file
View File

@@ -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으로 구현을 진행하면 됩니다.

198
README.md Normal file
View File

@@ -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

150
check_dependencies.py Normal file
View File

@@ -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())

122
install_dependencies.bat Normal file
View File

@@ -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

138
install_dependencies.ps1 Normal file
View File

@@ -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"

123
main.py Normal file
View File

@@ -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)

29
requirements.txt Normal file
View File

@@ -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

74
run.bat Normal file
View File

@@ -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

83
run.ps1 Normal file
View File

@@ -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"

5
src/__init__.py Normal file
View File

@@ -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"

33
src/connector/__init__.py Normal file
View File

@@ -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']

272
src/connector/config.py Normal file
View File

@@ -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}"
)

View File

@@ -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

7
src/server/__init__.py Normal file
View File

@@ -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']

605
src/server/handlers.py Normal file
View File

@@ -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)}"
)]

View File

@@ -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)}"
)]

178
src/server/mcp_server.py Normal file
View File

@@ -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"<base64 image data: {len(b64_data)} chars>"
# 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())

190
src/server/models.py Normal file
View File

@@ -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

55
src/utils/__init__.py Normal file
View File

@@ -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'
]

385
src/utils/image_utils.py Normal file
View File

@@ -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

387
src/utils/validation.py Normal file
View File

@@ -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

25
start.bat Normal file
View File

@@ -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

119
test_fixes.py Normal file
View File

@@ -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())

0
tests/__init__.py Normal file
View File

106
tests/run_tests.py Normal file
View File

@@ -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)

159
tests/test_config.py Normal file
View File

@@ -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()

308
tests/test_flux_client.py Normal file
View File

@@ -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()

360
tests/test_handlers.py Normal file
View File

@@ -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()

197
tests/test_image_utils.py Normal file
View File

@@ -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()

218
tests/test_validation.py Normal file
View File

@@ -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()

150
troubleshoot.bat Normal file
View File

@@ -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