import os import logging import requests import base64 from io import BytesIO from PIL import Image, ImageDraw, ImageFont import openai from typing import Optional, Dict, Any import time import random logger = logging.getLogger(__name__) class ImageGenerationError(Exception): """Custom exception for image generation failures""" pass def generate_image(prompt: str, content_type: str, content_info: str = "") -> Optional[str]: """ Main function to generate images for any D&D content type Args: prompt: Detailed image generation prompt content_type: Type of content (character_portrait, item, etc.) content_info: Brief info about content for placeholder Returns: str: Path to generated image or None if all methods fail """ # Check for API keys openai_key = os.getenv('OPENAI_API_KEY') stability_key = os.getenv('STABILITY_API_KEY') if openai_key: try: return generate_with_dalle(prompt) except Exception as e: logger.warning(f"DALL-E failed: {e}") if stability_key: try: return generate_with_stability(prompt) except Exception as e: logger.warning(f"Stability AI failed: {e}") # Generate placeholder if no APIs work return generate_placeholder(content_type, content_info) def generate_with_dalle(prompt: str, size: str = "1024x1024") -> Optional[str]: """Generate image using OpenAI's DALL-E""" try: logger.info(f"Generating image with DALL-E: {prompt[:50]}...") openai.api_key = os.getenv('OPENAI_API_KEY') response = openai.images.generate( model="dall-e-3", prompt=prompt, size=size, quality="standard", n=1, ) image_url = response.data[0].url logger.info("DALL-E image generated successfully") return image_url except Exception as e: logger.error(f"DALL-E generation failed: {e}") raise ImageGenerationError(f"DALL-E failed: {str(e)}") def generate_with_stability(prompt: str) -> Optional[str]: """Generate image using Stability AI""" try: logger.info(f"Generating image with Stability AI: {prompt[:50]}...") url = "https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image" headers = { "Accept": "application/json", "Content-Type": "application/json", "Authorization": f"Bearer {os.getenv('STABILITY_API_KEY')}", } body = { "text_prompts": [{"text": prompt, "weight": 1}], "cfg_scale": 7, "height": 1024, "width": 1024, "samples": 1, "steps": 30, } response = requests.post(url, headers=headers, json=body) if response.status_code != 200: raise ImageGenerationError(f"Stability API error: {response.status_code}") data = response.json() # Convert base64 to image image_data = base64.b64decode(data["artifacts"][0]["base64"]) img = Image.open(BytesIO(image_data)) # Save temporarily and return path temp_path = f"temp_{int(time.time())}_{random.randint(1000,9999)}.png" img.save(temp_path) logger.info("Stability AI image generated successfully") return temp_path except Exception as e: logger.error(f"Stability AI generation failed: {e}") raise ImageGenerationError(f"Stability AI failed: {str(e)}") def generate_placeholder(content_type: str, content_info: str) -> str: """Generate a themed placeholder image when AI services fail""" try: # Create different colored placeholders for different content types colors = { "character_portrait": (70, 130, 180), # Steel Blue "npc_portrait": (139, 69, 19), # Saddle Brown "item": (255, 215, 0), # Gold "location": (34, 139, 34), # Forest Green "faction_symbol": (128, 0, 128), # Purple "deity": (255, 255, 255), # White "scenario": (220, 20, 60) # Crimson } color = colors.get(content_type, (128, 128, 128)) # Create image img = Image.new('RGB', (512, 512), color=color) draw = ImageDraw.Draw(img) # Try to load a font, fall back to default if not available try: font = ImageFont.truetype("arial.ttf", 24) except: font = ImageFont.load_default() # Add text text_lines = [ f"{content_type.replace('_', ' ').title()}", "Placeholder Image", content_info[:30] + "..." if len(content_info) > 30 else content_info ] y_offset = 200 for line in text_lines: # Get text bounding box bbox = draw.textbbox((0, 0), line, font=font) text_width = bbox[2] - bbox[0] text_height = bbox[3] - bbox[1] # Center text x = (512 - text_width) // 2 draw.text((x, y_offset), line, fill=(255, 255, 255), font=font) y_offset += text_height + 10 # Save placeholder temp_path = f"placeholder_{content_type}_{int(time.time())}.png" img.save(temp_path) logger.info(f"Generated {content_type} placeholder image") return temp_path except Exception as e: logger.error(f"Failed to create placeholder: {e}") return None def cleanup_temp_files(): """Clean up temporary image files""" import glob temp_files = glob.glob("temp_*.png") + glob.glob("placeholder_*.png") for file in temp_files: try: os.remove(file) logger.info(f"Cleaned up temporary file: {file}") except Exception as e: logger.warning(f"Failed to clean up {file}: {e}")