Spaces:
Sleeping
Sleeping
File size: 6,199 Bytes
409063d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
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}") |