MCP Server - Images

"""Black Forest Labs image generation client.""" import httpx import asyncio import logging from pathlib import Path from typing import Optional, Dict, Any from enum import Enum logger = logging.getLogger(__name__) class BFLModel(str, Enum): """Valid models for Black Forest Labs API""" FLUX_PRO_11_ULTRA = "flux-pro-1.1-ultra" FLUX_PRO_11 = "flux-pro-1.1" FLUX_PRO = "flux-pro" FLUX_DEV = "flux-dev" class BFLClient: """Client for Black Forest Labs image generation.""" def __init__(self, api_key: str, base_url: str = "https://api.us1.bfl.ai/v1"): self.api_key = api_key self.base_url = base_url # Clean API key (remove prefix if present) api_key_clean = api_key.replace('bfl_', '') if api_key.startswith('bfl_') else api_key self.client = httpx.AsyncClient( base_url=base_url, headers={ "x-key": api_key_clean, "accept": "application/json", "Content-Type": "application/json" }, timeout=120.0 ) def _validate_model(self, model: str) -> str: """Validate and normalize model name.""" try: return BFLModel(model).value except ValueError: logger.warning(f"Invalid BFL model {model}, using default flux-pro-1.1") return BFLModel.FLUX_PRO_11.value def _convert_aspect_ratio_to_dimensions(self, aspect_ratio: str) -> tuple[int, int]: """Convert aspect ratio to width/height dimensions.""" # BFL uses pixel dimensions, not aspect ratio strings aspect_map = { "16:9": (1344, 768), "1:1": (1024, 1024), "21:9": (1536, 640), "2:3": (832, 1216), "3:2": (1216, 832), "4:5": (896, 1152), "5:4": (1152, 896), "9:16": (768, 1344), "9:21": (640, 1536), } return aspect_map.get(aspect_ratio, (1024, 1024)) async def generate_image( self, prompt: str, negative_prompt: Optional[str] = None, # Ignored - BFL doesn't support negative prompts model: str = "flux-pro-1.1", aspect_ratio: str = "1:1", cfg_scale: float = 7.0, # Ignored - BFL doesn't support CFG scale seed: Optional[int] = None, # Ignored - BFL doesn't support custom seeds output_path: Path = None ) -> Dict[str, Any]: """Generate image using Black Forest Labs API.""" # Validate and normalize model model = self._validate_model(model) # Convert aspect ratio to dimensions width, height = self._convert_aspect_ratio_to_dimensions(aspect_ratio) # Log ignored parameters if negative_prompt: logger.info("BFL doesn't support negative prompts - parameter ignored") if cfg_scale != 7.0: logger.info("BFL doesn't support cfg_scale - parameter ignored") if seed is not None: logger.info("BFL doesn't support custom seeds - parameter ignored") try: logger.info(f"Generating image with BFL: {model}") logger.debug(f"Prompt: {prompt[:100]}...") # Submit generation request response = await self.client.post( f"/{model}", json={ "prompt": prompt, "width": width, "height": height } ) # Handle specific BFL error codes if response.status_code == 402: raise RuntimeError("Insufficient BFL credits") elif response.status_code == 429: raise RuntimeError("Too many active BFL tasks") elif response.status_code != 200: try: error_json = response.json() error_msg = error_json.get('message', f'HTTP {response.status_code}') raise RuntimeError(f"BFL API Error: {error_msg}") except Exception: raise RuntimeError(f"BFL API Error: {response.status_code} - {response.text}") # Get request ID request_data = response.json() request_id = request_data["id"] logger.info(f"BFL generation started, ID: {request_id}") # Poll for results max_attempts = 120 # 2 minutes at 1 second intervals attempt = 0 while attempt < max_attempts: await asyncio.sleep(1.0) # Wait 1 second between polls attempt += 1 result_response = await self.client.get( "/get_result", params={"id": request_id} ) result_response.raise_for_status() result_data = result_response.json() status = result_data["status"] logger.debug(f"BFL generation status: {status} (attempt {attempt})") if status == "Ready": # Download image from signed URL image_url = result_data["result"]["sample"] async with httpx.AsyncClient() as dl_client: img_response = await dl_client.get(image_url) img_response.raise_for_status() # Save image if output path provided if output_path: output_path.parent.mkdir(parents=True, exist_ok=True) with open(output_path, "wb") as f: f.write(img_response.content) logger.info(f"Image saved to: {output_path}") # Return generation result return { "success": True, "provider": "bfl", "model": model, "image_size": len(img_response.content), "parameters": { "prompt": prompt, "model": model, "width": width, "height": height, "aspect_ratio": aspect_ratio, "output_format": "png" }, "request_id": request_id } elif status == "Failed": error_msg = result_data.get('error', 'Unknown error') raise RuntimeError(f"BFL generation failed: {error_msg}") # Timeout after max attempts raise RuntimeError(f"BFL generation timed out after {max_attempts} seconds") except httpx.TimeoutException: raise RuntimeError("Request timed out - BFL API took too long to respond") except httpx.HTTPError as e: raise RuntimeError(f"HTTP error: {str(e)}") except Exception as e: raise RuntimeError(f"Failed to generate image: {str(e)}") async def close(self): """Close the HTTP client.""" await self.client.aclose()