"""Image generation for Murder Mystery MCP Server.
Generates images using HuggingFace's Z-Image-Turbo model for:
- Scene/location images
- Suspect portraits
- Title cards
All images are styled to look like 1990s point-and-click adventure games.
"""
import hashlib
import logging
import os
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
logger = logging.getLogger(__name__)
# =============================================================================
# CONFIGURATION
# =============================================================================
HF_TOKEN = os.getenv("HF_TOKEN")
# Global HuggingFace client (lazy initialization)
_hf_client = None
# Cache directory
IMAGE_CACHE_DIR = Path(tempfile.gettempdir()) / "murder_mystery_images_mcp"
IMAGE_CACHE_DIR.mkdir(exist_ok=True)
def _get_cache_key(prompt: str) -> str:
"""Generate a cache key for a prompt."""
return hashlib.md5(prompt.encode()).hexdigest()[:12]
def _get_cached_image(cache_key: str) -> Optional[str]:
"""Check if image exists in cache and return path."""
cache_path = IMAGE_CACHE_DIR / f"{cache_key}.png"
if cache_path.exists():
logger.info(f"Found cached image: {cache_key}")
return str(cache_path)
return None
def get_hf_client():
"""Get HuggingFace Inference Client."""
global _hf_client
if _hf_client is None:
from huggingface_hub import InferenceClient
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
raise ValueError("HF_TOKEN environment variable not set")
_hf_client = InferenceClient(provider="fal-ai", api_key=hf_token) # type: ignore
logger.info("HuggingFace client initialized")
return _hf_client
# async def _download_and_cache_image(url: str, cache_key: str) -> Optional[str]:
# """Download image from URL and save to cache."""
# try:
# async with httpx.AsyncClient(timeout=30.0) as client:
# response = await client.get(url)
# if response.status_code == 200:
# cache_path = IMAGE_CACHE_DIR / f"{cache_key}.png"
# cache_path.write_bytes(response.content)
# logger.info(f"Cached image: {cache_path}")
# return str(cache_path)
# except Exception as e:
# logger.error(f"Failed to cache image: {e}")
# return None
# FAL_API_URL = "https://fal.run/fal-ai/flux/schnell"
# FAL_API_KEY = os.getenv("FAL_KEY")
# Art style suffix added to all prompts
ART_STYLE = """
Rendered in the distinctive style of 1990s LucasArts point-and-click adventure games
like Monkey Island, Day of the Tentacle, and Full Throttle. Hand-painted digital art
with visible brushwork and painterly textures. Rich, saturated color palette with bold
contrasts. Slightly stylized and exaggerated proportions, not photorealistic. Dramatic
chiaroscuro lighting with warm amber highlights and deep cool shadows.
"""
@dataclass
class ImageResult:
"""Result of image generation."""
url: Optional[str] = None
error: Optional[str] = None
generation_time: float = 0.0
prompt_used: str = ""
# =============================================================================
# IMAGE GENERATION FUNCTIONS
# =============================================================================
async def generate_image(prompt: str, width: int = 1024, height: int = 576) -> ImageResult:
"""Generate an image using HuggingFace Z-Image-Turbo model.
Args:
prompt: The image description
width: Image width (default 1024)
height: Image height (default 576)
Returns:
ImageResult with URL (local path) or error
"""
if not HF_TOKEN:
return ImageResult(error="HF_TOKEN environment variable not set")
# Add art style to prompt
full_prompt = f"{prompt}\n\n{ART_STYLE}"
# Check cache first
cache_key = _get_cache_key(full_prompt)
cached_path = _get_cached_image(cache_key)
if cached_path:
return ImageResult(
url=cached_path,
generation_time=0.0,
prompt_used=full_prompt
)
start_time = time.time()
try:
# Get HuggingFace client (lazy initialization)
client = get_hf_client()
# Generate image
image = client.text_to_image(
full_prompt,
model="Tongyi-MAI/Z-Image-Turbo",
width=width,
height=height,
)
# Save to cache
cache_path = IMAGE_CACHE_DIR / f"{cache_key}.png"
image.save(cache_path, "PNG")
logger.info(f"Cached image: {cache_path}")
return ImageResult(
url=str(cache_path),
generation_time=time.time() - start_time,
prompt_used=full_prompt
)
except ImportError:
return ImageResult(
error="huggingface_hub not installed. Run: pip install huggingface_hub",
generation_time=time.time() - start_time,
prompt_used=full_prompt
)
except ValueError as e:
return ImageResult(
error=str(e),
generation_time=time.time() - start_time,
prompt_used=full_prompt
)
except Exception as e:
logger.exception("Image generation failed")
return ImageResult(
error=str(e),
generation_time=time.time() - start_time,
prompt_used=full_prompt
)
# async def generate_image(prompt: str, width: int = 1024, height: int = 768) -> ImageResult:
# """Generate an image using fal.ai FLUX model.
# Args:
# prompt: The image description
# width: Image width (default 1024)
# height: Image height (default 768)
# Returns:
# ImageResult with URL or error
# """
# if not FAL_API_KEY:
# return ImageResult(error="FAL_KEY environment variable not set")
# # Add art style to prompt
# full_prompt = f"{prompt}\n\n{ART_STYLE}"
# # Check cache first
# cache_key = _get_cache_key(full_prompt)
# cached_path = _get_cached_image(cache_key)
# if cached_path:
# return ImageResult(
# url=cached_path, # Return local path instead of URL
# generation_time=0.0,
# prompt_used=full_prompt
# )
# start_time = time.time()
# try:
# async with httpx.AsyncClient(timeout=60.0) as client:
# response = await client.post(
# FAL_API_URL,
# headers={
# "Authorization": f"Key {FAL_API_KEY}",
# "Content-Type": "application/json"
# },
# json={
# "prompt": full_prompt,
# "image_size": {
# "width": width,
# "height": height
# },
# "num_inference_steps": 4, # FLUX schnell uses 4 steps
# "num_images": 1,
# "enable_safety_checker": True
# }
# )
# if response.status_code != 200:
# logger.error(f"FAL API error: {response.status_code} - {response.text}")
# return ImageResult(
# error=f"API error: {response.status_code}",
# generation_time=time.time() - start_time,
# prompt_used=full_prompt
# )
# data = response.json()
# images = data.get("images", [])
# if not images:
# return ImageResult(
# error="No images returned",
# generation_time=time.time() - start_time,
# prompt_used=full_prompt
# )
# image_url = images[0].get("url")
# # Download and cache the image
# cached_path = await _download_and_cache_image(image_url, cache_key)
# return ImageResult(
# url=cached_path or image_url, # Prefer cached path, fallback to URL
# generation_time=time.time() - start_time,
# prompt_used=full_prompt
# )
# except httpx.TimeoutException:
# return ImageResult(
# error="Request timed out",
# generation_time=time.time() - start_time,
# prompt_used=full_prompt
# )
# except Exception as e:
# logger.exception("Image generation failed")
# return ImageResult(
# error=str(e),
# generation_time=time.time() - start_time,
# prompt_used=full_prompt
# )
async def generate_scene(location_description: str, setting: str) -> ImageResult:
"""Generate a scene/location image.
Args:
location_description: Description of the location
setting: Overall mystery setting for context
Returns:
ImageResult with URL or error
"""
prompt = f"""A detailed scene from a murder mystery game showing: {location_description}
Setting context: {setting}
The scene should have atmospheric lighting suggesting mystery and intrigue.
Include environmental storytelling details that hint at the story.
Wide shot showing the full location."""
# Changed from 1024x768 to 1024x576 to match main code
return await generate_image(prompt, width=1024, height=576)
async def generate_portrait(character_description: str, character_name: str) -> ImageResult:
"""Generate a character portrait.
Args:
character_description: Description of the character
character_name: Character's name
Returns:
ImageResult with URL or error
"""
prompt = f"""Character portrait of {character_name} for a murder mystery game.
{character_description}
Head and shoulders portrait with dramatic three-quarter lighting.
Expressive face showing personality. Detailed costume appropriate to their role.
Background should be simple but atmospheric."""
# Changed from 768x768 to 1024x576 to match main code
return await generate_image(prompt, width=1024, height=576)
async def generate_title_card(description: str, victim_name: str) -> ImageResult:
"""Generate a title card for the mystery.
Args:
description: Description of the setting/scene
victim_name: Name of the victim
Returns:
ImageResult with URL or error
"""
prompt = f"""Dramatic title card for a murder mystery game.
{description}
Cinematic composition suggesting danger and intrigue.
Dark atmospheric mood with dramatic lighting.
Visual elements that hint at the crime without being graphic.
Space at the top or bottom for title text overlay."""
return await generate_image(prompt, width=1024, height=576) # 16:9 aspect
# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================
def is_image_generation_available() -> bool:
"""Check if image generation is available (HF_TOKEN is set)."""
return bool(HF_TOKEN)
async def test_image_generation() -> bool:
"""Test if image generation is working."""
if not HF_TOKEN:
return ImageResult(error="HF_TOKEN environment variable not set")
result = await generate_image("A simple test image", width=256, height=256)
return result.url is not None
