"""
Image handler implementations for Fal.ai MCP Server.
Contains: generate_image, generate_image_structured, generate_image_from_image
"""
import asyncio
import json
from typing import Any, Dict, List
from loguru import logger
from mcp.types import TextContent
from fal_mcp_server.model_registry import ModelRegistry
from fal_mcp_server.queue.base import QueueStrategy
async def handle_generate_image(
arguments: Dict[str, Any],
registry: ModelRegistry,
queue_strategy: QueueStrategy,
) -> List[TextContent]:
"""Handle the generate_image tool."""
model_input = arguments.get("model", "flux_schnell")
try:
model_id = await registry.resolve_model_id(model_input)
except ValueError as e:
return [
TextContent(
type="text",
text=f"❌ {e}. Use list_models to see available options.",
)
]
fal_args: Dict[str, Any] = {
"prompt": arguments["prompt"],
"image_size": arguments.get("image_size", "landscape_16_9"),
"num_images": arguments.get("num_images", 1),
}
# Add optional parameters
if "negative_prompt" in arguments:
fal_args["negative_prompt"] = arguments["negative_prompt"]
if "seed" in arguments:
fal_args["seed"] = arguments["seed"]
if "enable_safety_checker" in arguments:
fal_args["enable_safety_checker"] = arguments["enable_safety_checker"]
if "output_format" in arguments:
fal_args["output_format"] = arguments["output_format"]
# Use fast execution (no queue) for image generation
try:
result = await queue_strategy.execute_fast(model_id, fal_args)
except Exception as e:
logger.error("Image generation failed: %s", e)
return [
TextContent(
type="text",
text=f"❌ Image generation failed: {e}",
)
]
# Check for error in response
if "error" in result:
error_msg = result.get("error", "Unknown error")
logger.error("Image generation failed for %s: %s", model_id, error_msg)
return [
TextContent(
type="text",
text=f"❌ Image generation failed: {error_msg}",
)
]
images = result.get("images", [])
if not images:
logger.warning("Image generation returned no images. Model: %s", model_id)
return [
TextContent(
type="text",
text=f"❌ No images were generated by {model_id}. The prompt may have been filtered.",
)
]
# Extract URLs safely
try:
urls = [img["url"] for img in images]
except (KeyError, TypeError) as e:
logger.error("Malformed image response from %s: %s", model_id, e)
return [
TextContent(
type="text",
text=f"❌ Image generation completed but response was malformed: {e}",
)
]
response = f"🎨 Generated {len(urls)} image(s) with {model_id}:\n\n"
for i, url in enumerate(urls, 1):
response += f"Image {i}: {url}\n"
return [TextContent(type="text", text=response)]
async def handle_generate_image_structured(
arguments: Dict[str, Any],
registry: ModelRegistry,
queue_strategy: QueueStrategy,
) -> List[TextContent]:
"""Handle the generate_image_structured tool."""
model_input = arguments.get("model", "flux_schnell")
try:
model_id = await registry.resolve_model_id(model_input)
except ValueError as e:
return [
TextContent(
type="text",
text=f"❌ {e}. Use list_models to see available options.",
)
]
# Build structured JSON prompt from arguments
structured_prompt: Dict[str, Any] = {}
# Required field
structured_prompt["scene"] = arguments["scene"]
# Optional structured fields
for field in [
"subjects",
"style",
"color_palette",
"lighting",
"mood",
"background",
"composition",
"camera",
"effects",
]:
if field in arguments:
structured_prompt[field] = arguments[field]
# Convert structured prompt to JSON string
json_prompt = json.dumps(structured_prompt, indent=2)
fal_args: Dict[str, Any] = {
"prompt": json_prompt,
"image_size": arguments.get("image_size", "landscape_16_9"),
"num_images": arguments.get("num_images", 1),
}
# Add optional generation parameters
if "negative_prompt" in arguments:
fal_args["negative_prompt"] = arguments["negative_prompt"]
if "seed" in arguments:
fal_args["seed"] = arguments["seed"]
if "enable_safety_checker" in arguments:
fal_args["enable_safety_checker"] = arguments["enable_safety_checker"]
if "output_format" in arguments:
fal_args["output_format"] = arguments["output_format"]
# Use fast execution with timeout protection
logger.info("Starting structured image generation with %s", model_id)
try:
result = await asyncio.wait_for(
queue_strategy.execute_fast(model_id, fal_args),
timeout=60,
)
except asyncio.TimeoutError:
logger.error("Structured image generation timed out for %s", model_id)
return [
TextContent(
type="text",
text=f"❌ Image generation timed out after 60 seconds with {model_id}. Please try again.",
)
]
# Check for error in response
if "error" in result:
error_msg = result.get("error", "Unknown error")
logger.error(
"Structured image generation failed for %s: %s", model_id, error_msg
)
return [
TextContent(
type="text",
text=f"❌ Image generation failed: {error_msg}",
)
]
images = result.get("images", [])
if not images:
logger.warning(
"Structured image generation returned no images. Model: %s",
model_id,
)
return [
TextContent(
type="text",
text=f"❌ No images were generated by {model_id}. The prompt may have been filtered or the request format was invalid.",
)
]
# Extract URLs safely
try:
urls = [img["url"] for img in images]
except (KeyError, TypeError) as e:
logger.error("Malformed image response from %s: %s", model_id, e)
return [
TextContent(
type="text",
text=f"❌ Image generation completed but response was malformed: {e}",
)
]
response = (
f"🎨 Generated {len(urls)} image(s) with {model_id} (structured prompt):\n\n"
)
for i, url in enumerate(urls, 1):
response += f"Image {i}: {url}\n"
return [TextContent(type="text", text=response)]
async def handle_generate_image_from_image(
arguments: Dict[str, Any],
registry: ModelRegistry,
queue_strategy: QueueStrategy,
) -> List[TextContent]:
"""Handle the generate_image_from_image tool."""
model_input = arguments.get("model", "fal-ai/flux/dev/image-to-image")
try:
model_id = await registry.resolve_model_id(model_input)
except ValueError as e:
return [
TextContent(
type="text",
text=f"❌ {e}. Use list_models to see available options.",
)
]
# Both image_url and prompt are required
img2img_args: Dict[str, Any] = {
"image_url": arguments["image_url"],
"prompt": arguments["prompt"],
"strength": arguments.get("strength", 0.75),
"num_images": arguments.get("num_images", 1),
}
# Add optional parameters
if "negative_prompt" in arguments:
img2img_args["negative_prompt"] = arguments["negative_prompt"]
if "seed" in arguments:
img2img_args["seed"] = arguments["seed"]
if "enable_safety_checker" in arguments:
img2img_args["enable_safety_checker"] = arguments["enable_safety_checker"]
if "output_format" in arguments:
img2img_args["output_format"] = arguments["output_format"]
logger.info(
"Starting image-to-image transformation with %s from %s",
model_id,
(
arguments["image_url"][:50] + "..."
if len(arguments["image_url"]) > 50
else arguments["image_url"]
),
)
# Use fast execution with timeout protection
try:
result = await asyncio.wait_for(
queue_strategy.execute_fast(model_id, img2img_args),
timeout=60,
)
except asyncio.TimeoutError:
logger.error(
"Image-to-image transformation timed out after 60s. Model: %s",
model_id,
)
return [
TextContent(
type="text",
text=f"❌ Image transformation timed out after 60 seconds with {model_id}. Please try again.",
)
]
except Exception as e:
logger.exception("Image-to-image transformation failed: %s", e)
return [
TextContent(
type="text",
text=f"❌ Image transformation failed: {e}",
)
]
# Check for error in response
if "error" in result:
error_msg = result.get("error", "Unknown error")
logger.error(
"Image-to-image transformation failed for %s: %s",
model_id,
error_msg,
)
return [
TextContent(
type="text",
text=f"❌ Image transformation failed: {error_msg}",
)
]
images = result.get("images", [])
if not images:
logger.warning(
"Image-to-image transformation returned no images. Model: %s",
model_id,
)
return [
TextContent(
type="text",
text=f"❌ No images were generated by {model_id}. The source image may have been filtered.",
)
]
# Extract URLs safely
try:
urls = [img["url"] for img in images]
except (KeyError, TypeError) as e:
logger.error("Malformed image response from %s: %s", model_id, e)
return [
TextContent(
type="text",
text=f"❌ Image transformation completed but response was malformed: {e}",
)
]
response = f"🎨 Transformed image with {model_id}:\n\n"
response += f"**Source**: {arguments['image_url'][:50]}...\n\n"
for i, url in enumerate(urls, 1):
response += f"Result {i}: {url}\n"
return [TextContent(type="text", text=response)]
